Skip to content

Commit

Permalink
fix all mypy issues
Browse files Browse the repository at this point in the history
  • Loading branch information
@MartinBelthle
MartinBelthle committed Oct 1, 2024
1 parent 663e293 commit 534df03
Showing 1 changed file with 16 additions and 14 deletions.
30 changes: 16 additions & 14 deletions antarest/study/storage/rawstudy/model/filesystem/matrix/date_serializer.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,8 @@
#
# This file is part of the Antares project.

from __future__ import annotations

import re
from abc import ABC, abstractmethod
from typing import Hashable, List, Sequence, Tuple, cast
Expand Down Expand Up @@ -43,7 +45,7 @@ def __init__(self, area: str):
self.area = area

@abstractmethod
def extract_date(self, df: pd.DataFrame) -> Tuple[pd.Index, pd.DataFrame]:
def extract_date(self, df: pd.DataFrame) -> Tuple[pd.Index[str], pd.DataFrame]:
"""
Extract date from raw columns inside matrix file
Args:
Expand All @@ -54,7 +56,7 @@ def extract_date(self, df: pd.DataFrame) -> Tuple[pd.Index, pd.DataFrame]:
raise NotImplementedError()

@abstractmethod
def build_date(self, index: pd.Index) -> pd.DataFrame:
def build_date(self, index: pd.Index[str]) -> pd.DataFrame:
"""
Format in antares style date index
Args:
Expand All @@ -78,7 +80,7 @@ class HourlyMatrixSerializer(IDateMatrixSerializer):
Class implementation for hourly index
"""

def build_date(self, index: pd.Index) -> pd.DataFrame:
def build_date(self, index: pd.Index[str]) -> pd.DataFrame:
def _map(row: str) -> Tuple[str, int, str, str, str]:
m, d, h = re.split(r"[\s/]", row)
return "", 1, d, IDateMatrixSerializer._R_MONTHS[m], h
Expand All @@ -97,7 +99,7 @@ def _map(row: str) -> Tuple[str, int, str, str, str]:

return pd.concat([headers, matrix], axis=0)

def extract_date(self, df: pd.DataFrame) -> Tuple[pd.Index, pd.DataFrame]:
def extract_date(self, df: pd.DataFrame) -> Tuple[pd.Index[str], pd.DataFrame]:
# Extract left part with date
df_date = df.iloc[:, 2:5]
df_date.columns = pd.Index(data=["day", "month", "hour"])
Expand All @@ -108,15 +110,15 @@ def extract_date(self, df: pd.DataFrame) -> Tuple[pd.Index, pd.DataFrame]:
to_remove = cast(Sequence[Hashable], df.columns[0:5])
body = df.drop(to_remove, axis=1)

return pd.Index(date), body
return cast(pd.Index[str], pd.Index(date)), body


class DailyMatrixSerializer(IDateMatrixSerializer):
"""
Class implementation for daily index
"""

def build_date(self, index: pd.Index) -> pd.DataFrame:
def build_date(self, index: pd.Index[str]) -> pd.DataFrame:
def _map(row: str) -> Tuple[str, int, str, str]:
m, d = row.split("/")
return "", 1, d, IDateMatrixSerializer._R_MONTHS[m]
Expand All @@ -135,7 +137,7 @@ def _map(row: str) -> Tuple[str, int, str, str]:

return pd.concat([headers, matrix], axis=0)

def extract_date(self, df: pd.DataFrame) -> Tuple[pd.Index, pd.DataFrame]:
def extract_date(self, df: pd.DataFrame) -> Tuple[pd.Index[str], pd.DataFrame]:
# Extract left part with date
df_date = df.iloc[:, 2:4]
df_date.columns = pd.Index(["day", "month"])
Expand All @@ -146,15 +148,15 @@ def extract_date(self, df: pd.DataFrame) -> Tuple[pd.Index, pd.DataFrame]:
to_remove = cast(Sequence[Hashable], df.columns[0:4])
body = df.drop(to_remove, axis=1)

return pd.Index(date), body
return cast(pd.Index[str], pd.Index(date)), body


class WeeklyMatrixSerializer(IDateMatrixSerializer):
"""
Class implementation for weekly index
"""

def build_date(self, index: pd.Index) -> pd.DataFrame:
def build_date(self, index: pd.Index[str]) -> pd.DataFrame:
matrix = pd.DataFrame({0: [""] * index.size, 1: index.values})

headers = pd.DataFrame(
Expand All @@ -167,7 +169,7 @@ def build_date(self, index: pd.Index) -> pd.DataFrame:

return pd.concat([headers, matrix], axis=0)

def extract_date(self, df: pd.DataFrame) -> Tuple[pd.Index, pd.DataFrame]:
def extract_date(self, df: pd.DataFrame) -> Tuple[pd.Index[str], pd.DataFrame]:
# Extract left part with date
df_date = df.iloc[:, 1:2]
df_date.columns = pd.Index(["weekly"])
Expand All @@ -184,7 +186,7 @@ class MonthlyMatrixSerializer(IDateMatrixSerializer):
Class implementation for monthly index
"""

def build_date(self, index: pd.Index) -> pd.DataFrame:
def build_date(self, index: pd.Index[str]) -> pd.DataFrame:
matrix = pd.DataFrame(
{
0: [""] * index.size,
Expand All @@ -203,7 +205,7 @@ def build_date(self, index: pd.Index) -> pd.DataFrame:

return pd.concat([headers, matrix], axis=0)

def extract_date(self, df: pd.DataFrame) -> Tuple[pd.Index, pd.DataFrame]:
def extract_date(self, df: pd.DataFrame) -> Tuple[pd.Index[str], pd.DataFrame]:
# Extract left part with date
df_date = df.iloc[:, 2:3]
df_date.columns = pd.Index(["month"])
Expand All @@ -222,7 +224,7 @@ class AnnualMatrixSerializer(IDateMatrixSerializer):
Class implementation for annual index
"""

def build_date(self, index: pd.Index) -> pd.DataFrame:
def build_date(self, index: pd.Index[str]) -> pd.DataFrame:
return pd.DataFrame(
[
[self.area.upper(), "annual"],
Expand All @@ -232,7 +234,7 @@ def build_date(self, index: pd.Index) -> pd.DataFrame:
]
)

def extract_date(self, df: pd.DataFrame) -> Tuple[pd.Index, pd.DataFrame]:
def extract_date(self, df: pd.DataFrame) -> Tuple[pd.Index[str], pd.DataFrame]:
# Extract left part with date
df_date = df.iloc[:, 1:2]
df_date.columns = pd.Index(["annual"])
Expand Down

0 comments on commit 534df03

Please sign in to comment.