diff --git a/vessim/__init__.py b/vessim/__init__.py index 402b56e0..16afd1c4 100644 --- a/vessim/__init__.py +++ b/vessim/__init__.py @@ -4,7 +4,6 @@ from typing import Union, List, Optional, Any, Literal import pandas as pd -from pandas._typing import InterpolateOptions DatetimeLike = Union[str, datetime] @@ -136,6 +135,7 @@ def actual(self, dt: DatetimeLike, zone: Optional[str] = None) -> Any: ] # Mypy somehow has trouble with indexing in a dataframe with DatetimeIndex + # if dt in self._actual.index: return self._actual.loc[dt, zone] # type: ignore elif self._fill_method == "ffill" and dt > self._actual.index[0]: @@ -151,7 +151,7 @@ def forecast( end_time: DatetimeLike, zone: Optional[str] = None, frequency: Optional[Union[str, pd.DateOffset, timedelta]] = None, - resample_method: Optional[InterpolateOptions] = None, + resample_method: Optional[str] = None, ) -> pd.Series: """Retrieves of forecasted data points within window at a frequency. @@ -308,7 +308,7 @@ def _resample_to_frequency( start_time: datetime, end_time: datetime, frequency: pd.DateOffset, - resample_method: Optional[InterpolateOptions] = None, + resample_method: Optional[str] = None, ) -> pd.Series: """Transform series into the desired frequency between start and end time.""" # Cutoff data for performance @@ -329,7 +329,7 @@ def _resample_to_frequency( if resample_method == "ffill": df.ffill(inplace=True) else: - df.interpolate(method=resample_method, inplace=True) + df.interpolate(method=resample_method, inplace=True) # type: ignore # Get the data to the desired frequency after interpolation return df.loc[start_time:end_time].asfreq(frequency) # type: ignore