Skip to content

Commit b71e807

Browse files
authored
CLN: Add/refine type hints to some functions in core.dtypes.cast (#33286)
1 parent ad4ad22 commit b71e807

File tree

4 files changed

+28
-16
lines changed

4 files changed

+28
-16
lines changed

pandas/core/dtypes/cast.py

+13-7
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
"""
44

55
from datetime import date, datetime, timedelta
6+
from typing import TYPE_CHECKING, Type
67

78
import numpy as np
89

@@ -63,13 +64,18 @@
6364
ABCDataFrame,
6465
ABCDatetimeArray,
6566
ABCDatetimeIndex,
67+
ABCExtensionArray,
6668
ABCPeriodArray,
6769
ABCPeriodIndex,
6870
ABCSeries,
6971
)
7072
from pandas.core.dtypes.inference import is_list_like
7173
from pandas.core.dtypes.missing import isna, notna
7274

75+
if TYPE_CHECKING:
76+
from pandas import Series
77+
from pandas.core.arrays import ExtensionArray # noqa: F401
78+
7379
_int8_max = np.iinfo(np.int8).max
7480
_int16_max = np.iinfo(np.int16).max
7581
_int32_max = np.iinfo(np.int32).max
@@ -246,18 +252,16 @@ def trans(x):
246252
return result
247253

248254

249-
def maybe_cast_result(
250-
result, obj: ABCSeries, numeric_only: bool = False, how: str = ""
251-
):
255+
def maybe_cast_result(result, obj: "Series", numeric_only: bool = False, how: str = ""):
252256
"""
253257
Try casting result to a different type if appropriate
254258
255259
Parameters
256260
----------
257261
result : array-like
258262
Result to cast.
259-
obj : ABCSeries
260-
Input series from which result was calculated.
263+
obj : Series
264+
Input Series from which result was calculated.
261265
numeric_only : bool, default False
262266
Whether to cast only numerics or datetimes as well.
263267
how : str, default ""
@@ -313,13 +317,13 @@ def maybe_cast_result_dtype(dtype: DtypeObj, how: str) -> DtypeObj:
313317
return d.get((dtype, how), dtype)
314318

315319

316-
def maybe_cast_to_extension_array(cls, obj, dtype=None):
320+
def maybe_cast_to_extension_array(cls: Type["ExtensionArray"], obj, dtype=None):
317321
"""
318322
Call to `_from_sequence` that returns the object unchanged on Exception.
319323
320324
Parameters
321325
----------
322-
cls : ExtensionArray subclass
326+
cls : class, subclass of ExtensionArray
323327
obj : arraylike
324328
Values to pass to cls._from_sequence
325329
dtype : ExtensionDtype, optional
@@ -329,6 +333,8 @@ def maybe_cast_to_extension_array(cls, obj, dtype=None):
329333
ExtensionArray or obj
330334
"""
331335
assert isinstance(cls, type), f"must pass a type: {cls}"
336+
assertion_msg = f"must pass a subclass of ExtensionArray: {cls}"
337+
assert issubclass(cls, ABCExtensionArray), assertion_msg
332338
try:
333339
result = cls._from_sequence(obj, dtype=dtype)
334340
except Exception:

pandas/core/groupby/generic.py

+3-3
Original file line numberDiff line numberDiff line change
@@ -151,7 +151,7 @@ def pinner(cls):
151151

152152

153153
@pin_whitelisted_properties(Series, base.series_apply_whitelist)
154-
class SeriesGroupBy(GroupBy):
154+
class SeriesGroupBy(GroupBy[Series]):
155155
_apply_whitelist = base.series_apply_whitelist
156156

157157
def _iterate_slices(self) -> Iterable[Series]:
@@ -815,7 +815,7 @@ def pct_change(self, periods=1, fill_method="pad", limit=None, freq=None):
815815

816816

817817
@pin_whitelisted_properties(DataFrame, base.dataframe_apply_whitelist)
818-
class DataFrameGroupBy(GroupBy):
818+
class DataFrameGroupBy(GroupBy[DataFrame]):
819819

820820
_apply_whitelist = base.dataframe_apply_whitelist
821821

@@ -1462,7 +1462,7 @@ def _transform_fast(self, result: DataFrame, func_nm: str) -> DataFrame:
14621462
for i, _ in enumerate(result.columns):
14631463
res = algorithms.take_1d(result.iloc[:, i].values, ids)
14641464
# TODO: we have no test cases that get here with EA dtypes;
1465-
# try_cast may not be needed if EAs never get here
1465+
# maybe_cast_result may not be needed if EAs never get here
14661466
if cast:
14671467
res = maybe_cast_result(res, obj.iloc[:, i], how=func_nm)
14681468
output.append(res)

pandas/core/groupby/groupby.py

+11-5
Original file line numberDiff line numberDiff line change
@@ -17,13 +17,15 @@ class providing the base-class of operations.
1717
Callable,
1818
Dict,
1919
FrozenSet,
20+
Generic,
2021
Hashable,
2122
Iterable,
2223
List,
2324
Mapping,
2425
Optional,
2526
Tuple,
2627
Type,
28+
TypeVar,
2729
Union,
2830
)
2931

@@ -353,13 +355,13 @@ def _group_selection_context(groupby):
353355
]
354356

355357

356-
class _GroupBy(PandasObject, SelectionMixin):
358+
class _GroupBy(PandasObject, SelectionMixin, Generic[FrameOrSeries]):
357359
_group_selection = None
358360
_apply_whitelist: FrozenSet[str] = frozenset()
359361

360362
def __init__(
361363
self,
362-
obj: NDFrame,
364+
obj: FrameOrSeries,
363365
keys: Optional[_KeysArgType] = None,
364366
axis: int = 0,
365367
level=None,
@@ -995,7 +997,11 @@ def _apply_filter(self, indices, dropna):
995997
return filtered
996998

997999

998-
class GroupBy(_GroupBy):
1000+
# To track operations that expand dimensions, like ohlc
1001+
OutputFrameOrSeries = TypeVar("OutputFrameOrSeries", bound=NDFrame)
1002+
1003+
1004+
class GroupBy(_GroupBy[FrameOrSeries]):
9991005
"""
10001006
Class for grouping and aggregating relational data.
10011007
@@ -2420,8 +2426,8 @@ def tail(self, n=5):
24202426
return self._selected_obj[mask]
24212427

24222428
def _reindex_output(
2423-
self, output: FrameOrSeries, fill_value: Scalar = np.NaN
2424-
) -> FrameOrSeries:
2429+
self, output: OutputFrameOrSeries, fill_value: Scalar = np.NaN
2430+
) -> OutputFrameOrSeries:
24252431
"""
24262432
If we have categorical groupers, then we might want to make sure that
24272433
we have a fully re-indexed output to the levels. This means expanding

pandas/core/groupby/ops.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -682,7 +682,7 @@ def _aggregate_series_pure_python(self, obj: Series, func):
682682

683683
assert result is not None
684684
result = lib.maybe_convert_objects(result, try_float=0)
685-
# TODO: try_cast back to EA?
685+
# TODO: maybe_cast_to_extension_array?
686686

687687
return result, counts
688688

0 commit comments

Comments
 (0)