Clean up typing of function signatures

ludwig/data/dataframe/daft.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Callable, Dict, List
+from typing import Callable, Dict, List, overload, Union
 
 import daft
 import pandas as pd
@@ -11,6 +11,10 @@
 logger = logging.getLogger(__name__)
 
 
+DataFrameMapFn = Callable[[pd.DataFrame], pd.DataFrame]
+SeriesMapFn = Callable[[pd.Series], pd.Series]
+
+
 class LudwigDaftDataframe:
     """Shim layer on top of a daft.Dataframe to make it behave more like a Pandas Dataframe object.
 
@@ -129,9 +133,22 @@ def map_objects(self, series: LudwigDaftSeries, map_fn: Callable[[object], objec
         # can be much more optimized in terms of memory usage
         return LudwigDaftSeries(series.expr.apply(map_fn, return_dtype=daft.DataType.python()))
 
-    def map_partitions(self, obj: LudwigDaftSeries, map_fn: Callable[[pd.Series], pd.Series], meta=None):
-        # NOTE: Although the function signature indicates that this function takes in a Series, in practice
-        # it appears that this function is often used interchangeably to run on both Series and DataFrames
+    # NOTE: Although the base class' function signature indicates that this function takes in a Series, in practice
+    # it appears that this function is often used interchangeably to run functions on both Series and DataFrames
+    @overload
+    def map_partitions(self, obj: LudwigDaftDataframe, map_fn: DataFrameMapFn):
+        ...
+
+    @overload
+    def map_partitions(self, obj: LudwigDaftSeries, map_fn: SeriesMapFn, meta=None):
+        ...
+
+    def map_partitions(
+        self,
+        obj: Union[LudwigDaftSeries, LudwigDaftDataframe],
+        map_fn: Union[DataFrameMapFn, SeriesMapFn],
+        meta=None,
+    ):
         if isinstance(obj, LudwigDaftDataframe):
             raise NotImplementedError("TODO: Implementation")
         elif isinstance(obj, LudwigDaftSeries):