CLN: Add/refine type hints to some functions in core.dtypes.cast (#33286

)

pandas/core/dtypes/cast.py

@@ -3,6 +3,7 @@
 """
 
 from datetime import date, datetime, timedelta
+from typing import TYPE_CHECKING, Type
 
 import numpy as np
 
@@ -63,13 +64,18 @@
     ABCDataFrame,
     ABCDatetimeArray,
     ABCDatetimeIndex,
+    ABCExtensionArray,
     ABCPeriodArray,
     ABCPeriodIndex,
     ABCSeries,
 )
 from pandas.core.dtypes.inference import is_list_like
 from pandas.core.dtypes.missing import isna, notna
 
+if TYPE_CHECKING:
+    from pandas import Series
+    from pandas.core.arrays import ExtensionArray  # noqa: F401
+
 _int8_max = np.iinfo(np.int8).max
 _int16_max = np.iinfo(np.int16).max
 _int32_max = np.iinfo(np.int32).max
@@ -246,18 +252,16 @@ def trans(x):
     return result
 
 
-def maybe_cast_result(
-    result, obj: ABCSeries, numeric_only: bool = False, how: str = ""
-):
+def maybe_cast_result(result, obj: "Series", numeric_only: bool = False, how: str = ""):
     """
     Try casting result to a different type if appropriate
 
     Parameters
     ----------
     result : array-like
         Result to cast.
-    obj : ABCSeries
-        Input series from which result was calculated.
+    obj : Series
+        Input Series from which result was calculated.
     numeric_only : bool, default False
         Whether to cast only numerics or datetimes as well.
     how : str, default ""
@@ -313,13 +317,13 @@ def maybe_cast_result_dtype(dtype: DtypeObj, how: str) -> DtypeObj:
     return d.get((dtype, how), dtype)
 
 
-def maybe_cast_to_extension_array(cls, obj, dtype=None):
+def maybe_cast_to_extension_array(cls: Type["ExtensionArray"], obj, dtype=None):
     """
     Call to `_from_sequence` that returns the object unchanged on Exception.
 
     Parameters
     ----------
-    cls : ExtensionArray subclass
+    cls : class, subclass of ExtensionArray
     obj : arraylike
         Values to pass to cls._from_sequence
     dtype : ExtensionDtype, optional
@@ -329,6 +333,8 @@ def maybe_cast_to_extension_array(cls, obj, dtype=None):
     ExtensionArray or obj
     """
     assert isinstance(cls, type), f"must pass a type: {cls}"
+    assertion_msg = f"must pass a subclass of ExtensionArray: {cls}"
+    assert issubclass(cls, ABCExtensionArray), assertion_msg
     try:
         result = cls._from_sequence(obj, dtype=dtype)
     except Exception:

pandas/core/groupby/generic.py

@@ -151,7 +151,7 @@ def pinner(cls):
 
 
 @pin_whitelisted_properties(Series, base.series_apply_whitelist)
-class SeriesGroupBy(GroupBy):
+class SeriesGroupBy(GroupBy[Series]):
     _apply_whitelist = base.series_apply_whitelist
 
     def _iterate_slices(self) -> Iterable[Series]:
@@ -815,7 +815,7 @@ def pct_change(self, periods=1, fill_method="pad", limit=None, freq=None):
 
 
 @pin_whitelisted_properties(DataFrame, base.dataframe_apply_whitelist)
-class DataFrameGroupBy(GroupBy):
+class DataFrameGroupBy(GroupBy[DataFrame]):
 
     _apply_whitelist = base.dataframe_apply_whitelist
 
@@ -1462,7 +1462,7 @@ def _transform_fast(self, result: DataFrame, func_nm: str) -> DataFrame:
         for i, _ in enumerate(result.columns):
             res = algorithms.take_1d(result.iloc[:, i].values, ids)
             # TODO: we have no test cases that get here with EA dtypes;
-            #  try_cast may not be needed if EAs never get here
+            #  maybe_cast_result may not be needed if EAs never get here
             if cast:
                 res = maybe_cast_result(res, obj.iloc[:, i], how=func_nm)
             output.append(res)

pandas/core/groupby/groupby.py

@@ -17,13 +17,15 @@ class providing the base-class of operations.
     Callable,
     Dict,
     FrozenSet,
+    Generic,
     Hashable,
     Iterable,
     List,
     Mapping,
     Optional,
     Tuple,
     Type,
+    TypeVar,
     Union,
 )
 
@@ -353,13 +355,13 @@ def _group_selection_context(groupby):
 ]
 
 
-class _GroupBy(PandasObject, SelectionMixin):
+class _GroupBy(PandasObject, SelectionMixin, Generic[FrameOrSeries]):
     _group_selection = None
     _apply_whitelist: FrozenSet[str] = frozenset()
 
     def __init__(
         self,
-        obj: NDFrame,
+        obj: FrameOrSeries,
         keys: Optional[_KeysArgType] = None,
         axis: int = 0,
         level=None,
@@ -995,7 +997,11 @@ def _apply_filter(self, indices, dropna):
         return filtered
 
 
-class GroupBy(_GroupBy):
+# To track operations that expand dimensions, like ohlc
+OutputFrameOrSeries = TypeVar("OutputFrameOrSeries", bound=NDFrame)
+
+
+class GroupBy(_GroupBy[FrameOrSeries]):
     """
     Class for grouping and aggregating relational data.
 
@@ -2420,8 +2426,8 @@ def tail(self, n=5):
         return self._selected_obj[mask]
 
     def _reindex_output(
-        self, output: FrameOrSeries, fill_value: Scalar = np.NaN
-    ) -> FrameOrSeries:
+        self, output: OutputFrameOrSeries, fill_value: Scalar = np.NaN
+    ) -> OutputFrameOrSeries:
         """
         If we have categorical groupers, then we might want to make sure that
         we have a fully re-indexed output to the levels. This means expanding

pandas/core/groupby/ops.py

@@ -682,7 +682,7 @@ def _aggregate_series_pure_python(self, obj: Series, func):
 
         assert result is not None
         result = lib.maybe_convert_objects(result, try_float=0)
-        # TODO: try_cast back to EA?
+        # TODO: maybe_cast_to_extension_array?
 
         return result, counts