mypy fixes

mypy fix
NREL · Dec 22, 2024 · 931dc9c · 931dc9c
1 parent 8feb0eb
commit 931dc9c
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 10 deletions.
diff --git a/src/chronify/sqlalchemy/functions.py b/src/chronify/sqlalchemy/functions.py
@@ -4,7 +4,7 @@
 in memory.
 """
 
-from typing import Any, Literal, TypeAlias
+from typing import Any, Literal, TypeAlias, Sequence
 
 import pandas as pd
 import polars as pl
@@ -44,7 +44,7 @@ def write_database(
     df: pd.DataFrame,
     conn: Connection,
     table_name: str,
-    configs: list[TimeBaseModel],
+    configs: Sequence[TimeBaseModel],
     if_table_exists: DbWriteMode = "append",
 ) -> None:
     """Write a Pandas DataFrame to the database.

diff --git a/src/chronify/time_utils.py b/src/chronify/time_utils.py
@@ -1,8 +1,8 @@
 """Functions related to time"""
 
 import logging
-
 import numpy as np
+from typing import Any
 
 import pandas as pd
 
@@ -15,8 +15,10 @@
 
 
 def shift_time_interval(
-    df: pd.Series, from_interval_type: TimeIntervalType, to_interval_type: TimeIntervalType
-) -> pd.Series:
+    df: pd.Series[pd.Timestamp],
+    from_interval_type: TimeIntervalType,
+    to_interval_type: TimeIntervalType,
+) -> Any:
     """Shift pandas timeseries by time interval based on interval type."""
     assert (
         from_interval_type != to_interval_type
@@ -40,8 +42,10 @@ def shift_time_interval(
 
 
 def roll_time_interval(
-    df: pd.Series, from_interval_type: TimeIntervalType, to_interval_type: TimeIntervalType
-) -> pd.Series:
+    df: pd.Series[pd.Timestamp],
+    from_interval_type: TimeIntervalType,
+    to_interval_type: TimeIntervalType,
+) -> np.typing.NDArray[Any]:
     """Roll pandas timeseries by time interval based on interval type with np.roll(),
     which includes time-wrapping.
     """
@@ -61,9 +65,11 @@ def roll_time_interval(
     return np.roll(df, shift)
 
 
-def wrap_timestamps(df: pd.Series, to_timestamps: list[pd.Timestamp]) -> pd.Series:
+def wrap_timestamps(
+    df: pd.Series[pd.Timestamp], to_timestamps: list[pd.Timestamp]
+) -> pd.Series[pd.Timestamp]:
     """Wrap pandas timeseries so it conforms to a list of timestamps."""
-    arr = np.sort(to_timestamps)
+    arr = np.sort(np.array(to_timestamps))
     freqs = set((np.roll(arr, -1) - arr)[:-1])
     assert len(freqs), f"Timeseries has more than one frequency, {freqs}"
     freq = list(freqs)[0]

diff --git a/tests/test_mapper_datetime_to_datetime.py b/tests/test_mapper_datetime_to_datetime.py
@@ -132,7 +132,6 @@ def test_time_interval_shift(
     )
     df = generate_datetime_dataframe(from_schema)
     to_schema = get_datetime_schema(2020, tzinfo, TimeIntervalType.PERIOD_ENDING, "to_table")
-
     error = ()
     run_test(iter_engines, df, from_schema, to_schema, error)