Fix Pandas type hints

spetlr-org · May 12, 2024 · 5d18056 · 5d18056
1 parent f895542
commit 5d18056
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/src/spetlr/power_bi/SparkPandasDataFrame.py b/src/spetlr/power_bi/SparkPandasDataFrame.py
@@ -4,7 +4,8 @@
 import numpy as np
 import pandas as pd
 from dateutil.parser import parse
-from pandas.core.generic import NDFrameT
+from pandas.core.frame import DataFrame as PandasDataFrame
+from pandas.core.generic import NDFrameT as PandasColumn
 from pyspark.sql import DataFrame
 from pytz import timezone, utc
 
@@ -26,7 +27,9 @@ class SparkPandasDataFrame:
     def __init__(
         self,
         json: Union[Dict, List],
-        schema: List[tuple[Union[str, Callable[[pd.DataFrame], NDFrameT]], str, str]],
+        schema: List[
+            tuple[Union[str, Callable[[PandasDataFrame], PandasColumn]], str, str]
+        ],
         *,
         indexing_columns: Union[int, List[int], str, List[str], None] = None,
         sorting_columns: Union[int, List[int], str, List[str], None] = None,
@@ -193,7 +196,7 @@ def _localize_time(
         time = time.astimezone(timezone(local_timezone_name))
         return time.replace(tzinfo=None)
 
-    def get_pandas_df(self) -> Union[pd.DataFrame, None]:
+    def get_pandas_df(self) -> Union[PandasDataFrame, None]:
         """
         Returns the data frame as a Pandas data frame.