From b4699fd9f30bd8b2e3223a3cc3f694a4afb710ae Mon Sep 17 00:00:00 2001 From: Brock Date: Sat, 8 May 2021 11:36:16 -0700 Subject: [PATCH 1/2] TYP: SelectionMixin --- pandas/core/base.py | 51 +++++++++++----------------------- pandas/core/groupby/groupby.py | 8 ++++-- pandas/core/groupby/grouper.py | 10 +++---- pandas/core/window/rolling.py | 7 +++-- 4 files changed, 31 insertions(+), 45 deletions(-) diff --git a/pandas/core/base.py b/pandas/core/base.py index 3270e3dd82f7d..3d6b2d678c735 100644 --- a/pandas/core/base.py +++ b/pandas/core/base.py @@ -8,6 +8,7 @@ from typing import ( TYPE_CHECKING, Any, + Hashable, TypeVar, cast, ) @@ -19,6 +20,7 @@ ArrayLike, Dtype, DtypeObj, + FrameOrSeriesUnion, IndexLabel, Shape, final, @@ -169,7 +171,9 @@ class SelectionMixin: object sub-classes need to define: obj, exclusions """ + obj: FrameOrSeriesUnion _selection: IndexLabel | None = None + exclusions: frozenset[Hashable] _internal_names = ["_cache", "__setstate__"] _internal_names_set = set(_internal_names) @@ -194,15 +198,10 @@ def _selection_list(self): @cache_readonly def _selected_obj(self): - # error: "SelectionMixin" has no attribute "obj" - if self._selection is None or isinstance( - self.obj, ABCSeries # type: ignore[attr-defined] - ): - # error: "SelectionMixin" has no attribute "obj" - return self.obj # type: ignore[attr-defined] + if self._selection is None or isinstance(self.obj, ABCSeries): + return self.obj else: - # error: "SelectionMixin" has no attribute "obj" - return self.obj[self._selection] # type: ignore[attr-defined] + return self.obj[self._selection] @cache_readonly def ndim(self) -> int: @@ -211,49 +210,31 @@ def ndim(self) -> int: @final @cache_readonly def _obj_with_exclusions(self): - # error: "SelectionMixin" has no attribute "obj" - if self._selection is not None and isinstance( - self.obj, ABCDataFrame # type: ignore[attr-defined] - ): - # error: "SelectionMixin" has no attribute "obj" - return self.obj.reindex( # type: ignore[attr-defined] - columns=self._selection_list - ) + if self._selection is not None and isinstance(self.obj, ABCDataFrame): + return self.obj.reindex(columns=self._selection_list) - # error: "SelectionMixin" has no attribute "exclusions" - if len(self.exclusions) > 0: # type: ignore[attr-defined] - # error: "SelectionMixin" has no attribute "obj" - # error: "SelectionMixin" has no attribute "exclusions" - return self.obj.drop(self.exclusions, axis=1) # type: ignore[attr-defined] + if len(self.exclusions) > 0: + return self.obj.drop(self.exclusions, axis=1) else: - # error: "SelectionMixin" has no attribute "obj" - return self.obj # type: ignore[attr-defined] + return self.obj def __getitem__(self, key): if self._selection is not None: raise IndexError(f"Column(s) {self._selection} already selected") if isinstance(key, (list, tuple, ABCSeries, ABCIndex, np.ndarray)): - # error: "SelectionMixin" has no attribute "obj" - if len( - self.obj.columns.intersection(key) # type: ignore[attr-defined] - ) != len(key): - # error: "SelectionMixin" has no attribute "obj" - bad_keys = list( - set(key).difference(self.obj.columns) # type: ignore[attr-defined] - ) + if len(self.obj.columns.intersection(key)) != len(key): + bad_keys = list(set(key).difference(self.obj.columns)) raise KeyError(f"Columns not found: {str(bad_keys)[1:-1]}") return self._gotitem(list(key), ndim=2) elif not getattr(self, "as_index", False): - # error: "SelectionMixin" has no attribute "obj" - if key not in self.obj.columns: # type: ignore[attr-defined] + if key not in self.obj.columns: raise KeyError(f"Column not found: {key}") return self._gotitem(key, ndim=2) else: - # error: "SelectionMixin" has no attribute "obj" - if key not in self.obj: # type: ignore[attr-defined] + if key not in self.obj: raise KeyError(f"Column not found: {key}") return self._gotitem(key, ndim=1) diff --git a/pandas/core/groupby/groupby.py b/pandas/core/groupby/groupby.py index 1105c1bd1d782..b12b6ba53bf39 100644 --- a/pandas/core/groupby/groupby.py +++ b/pandas/core/groupby/groupby.py @@ -588,7 +588,9 @@ class BaseGroupBy(PandasObject, SelectionMixin, Generic[FrameOrSeries]): axis: int grouper: ops.BaseGrouper - obj: FrameOrSeries + # error: Incompatible types in assignment (expression has type "FrameOrSeries", + # base class "SelectionMixin" defined the type as "Union[DataFrame, Series]") + obj: FrameOrSeries # type: ignore[assignment] group_keys: bool @final @@ -852,7 +854,7 @@ def __init__( axis: int = 0, level: IndexLabel | None = None, grouper: ops.BaseGrouper | None = None, - exclusions: set[Hashable] | None = None, + exclusions: frozenset[Hashable] | None = None, selection: IndexLabel | None = None, as_index: bool = True, sort: bool = True, @@ -901,7 +903,7 @@ def __init__( self.obj = obj self.axis = obj._get_axis_number(axis) self.grouper = grouper - self.exclusions = exclusions or set() + self.exclusions = frozenset(exclusions) if exclusions else frozenset() def __getattr__(self, attr: str): if attr in self._internal_names_set: diff --git a/pandas/core/groupby/grouper.py b/pandas/core/groupby/grouper.py index f1762a2535ff7..4650dbea27de1 100644 --- a/pandas/core/groupby/grouper.py +++ b/pandas/core/groupby/grouper.py @@ -652,7 +652,7 @@ def get_grouper( mutated: bool = False, validate: bool = True, dropna: bool = True, -) -> tuple[ops.BaseGrouper, set[Hashable], FrameOrSeries]: +) -> tuple[ops.BaseGrouper, frozenset[Hashable], FrameOrSeries]: """ Create and return a BaseGrouper, which is an internal mapping of how to create the grouper indexers. @@ -728,13 +728,13 @@ def get_grouper( if isinstance(key, Grouper): binner, grouper, obj = key._get_grouper(obj, validate=False) if key.key is None: - return grouper, set(), obj + return grouper, frozenset(), obj else: - return grouper, {key.key}, obj + return grouper, frozenset({key.key}), obj # already have a BaseGrouper, just return it elif isinstance(key, ops.BaseGrouper): - return key, set(), obj + return key, frozenset(), obj if not isinstance(key, list): keys = [key] @@ -861,7 +861,7 @@ def is_in_obj(gpr) -> bool: grouper = ops.BaseGrouper( group_axis, groupings, sort=sort, mutated=mutated, dropna=dropna ) - return grouper, exclusions, obj + return grouper, frozenset(exclusions), obj def _is_label_like(val) -> bool: diff --git a/pandas/core/window/rolling.py b/pandas/core/window/rolling.py index b51875134c614..f37c29414e479 100644 --- a/pandas/core/window/rolling.py +++ b/pandas/core/window/rolling.py @@ -13,6 +13,7 @@ TYPE_CHECKING, Any, Callable, + Hashable, ) import warnings @@ -109,7 +110,7 @@ class BaseWindow(SelectionMixin): """Provides utilities for performing windowing operations.""" _attributes: list[str] = [] - exclusions: set[str] = set() + exclusions: frozenset[Hashable] = frozenset() def __init__( self, @@ -123,7 +124,9 @@ def __init__( closed: str | None = None, method: str = "single", ): - self.obj = obj + # error: Incompatible types in assignment (expression has type "FrameOrSeries", + # variable has type "Union[DataFrame, Series]") + self.obj = obj # type: ignore[assignment] self.on = on self.closed = closed self.window = window From 7eb31809fb6817e51530a9a5cc2c7141100c4c64 Mon Sep 17 00:00:00 2001 From: Simon Hawkins Date: Sat, 8 May 2021 20:50:33 +0100 Subject: [PATCH 2/2] remove ignores --- pandas/_typing.py | 2 ++ pandas/core/apply.py | 10 +++------- pandas/core/base.py | 7 ++++--- pandas/core/groupby/groupby.py | 7 +------ pandas/core/window/rolling.py | 4 +--- 5 files changed, 11 insertions(+), 19 deletions(-) diff --git a/pandas/_typing.py b/pandas/_typing.py index 1e1fffdd60676..7763b0ceb610a 100644 --- a/pandas/_typing.py +++ b/pandas/_typing.py @@ -56,6 +56,7 @@ from pandas.core.generic import NDFrame from pandas.core.groupby.generic import ( DataFrameGroupBy, + GroupBy, SeriesGroupBy, ) from pandas.core.indexes.base import Index @@ -158,6 +159,7 @@ AggObjType = Union[ "Series", "DataFrame", + "GroupBy", "SeriesGroupBy", "DataFrameGroupBy", "BaseWindow", diff --git a/pandas/core/apply.py b/pandas/core/apply.py index ad25eb6fbcaa8..9d3437fe08b24 100644 --- a/pandas/core/apply.py +++ b/pandas/core/apply.py @@ -24,6 +24,7 @@ AggFuncTypeDict, AggObjType, Axis, + FrameOrSeries, FrameOrSeriesUnion, ) from pandas.util._decorators import cache_readonly @@ -60,10 +61,7 @@ Index, Series, ) - from pandas.core.groupby import ( - DataFrameGroupBy, - SeriesGroupBy, - ) + from pandas.core.groupby import GroupBy from pandas.core.resample import Resampler from pandas.core.window.rolling import BaseWindow @@ -1089,11 +1087,9 @@ def apply_standard(self) -> FrameOrSeriesUnion: class GroupByApply(Apply): - obj: SeriesGroupBy | DataFrameGroupBy - def __init__( self, - obj: SeriesGroupBy | DataFrameGroupBy, + obj: GroupBy[FrameOrSeries], func: AggFuncType, args, kwargs, diff --git a/pandas/core/base.py b/pandas/core/base.py index 3d6b2d678c735..542fd54ce0ac7 100644 --- a/pandas/core/base.py +++ b/pandas/core/base.py @@ -8,6 +8,7 @@ from typing import ( TYPE_CHECKING, Any, + Generic, Hashable, TypeVar, cast, @@ -20,7 +21,7 @@ ArrayLike, Dtype, DtypeObj, - FrameOrSeriesUnion, + FrameOrSeries, IndexLabel, Shape, final, @@ -165,13 +166,13 @@ class SpecificationError(Exception): pass -class SelectionMixin: +class SelectionMixin(Generic[FrameOrSeries]): """ mixin implementing the selection & aggregation interface on a group-like object sub-classes need to define: obj, exclusions """ - obj: FrameOrSeriesUnion + obj: FrameOrSeries _selection: IndexLabel | None = None exclusions: frozenset[Hashable] _internal_names = ["_cache", "__setstate__"] diff --git a/pandas/core/groupby/groupby.py b/pandas/core/groupby/groupby.py index b12b6ba53bf39..0a176dafb34bb 100644 --- a/pandas/core/groupby/groupby.py +++ b/pandas/core/groupby/groupby.py @@ -20,7 +20,6 @@ class providing the base-class of operations. from typing import ( TYPE_CHECKING, Callable, - Generic, Hashable, Iterable, Iterator, @@ -567,7 +566,7 @@ def group_selection_context(groupby: GroupBy) -> Iterator[GroupBy]: ] -class BaseGroupBy(PandasObject, SelectionMixin, Generic[FrameOrSeries]): +class BaseGroupBy(PandasObject, SelectionMixin[FrameOrSeries]): _group_selection: IndexLabel | None = None _apply_allowlist: frozenset[str] = frozenset() _hidden_attrs = PandasObject._hidden_attrs | { @@ -588,9 +587,6 @@ class BaseGroupBy(PandasObject, SelectionMixin, Generic[FrameOrSeries]): axis: int grouper: ops.BaseGrouper - # error: Incompatible types in assignment (expression has type "FrameOrSeries", - # base class "SelectionMixin" defined the type as "Union[DataFrame, Series]") - obj: FrameOrSeries # type: ignore[assignment] group_keys: bool @final @@ -842,7 +838,6 @@ class GroupBy(BaseGroupBy[FrameOrSeries]): more """ - obj: FrameOrSeries grouper: ops.BaseGrouper as_index: bool diff --git a/pandas/core/window/rolling.py b/pandas/core/window/rolling.py index f37c29414e479..490cdca8519e6 100644 --- a/pandas/core/window/rolling.py +++ b/pandas/core/window/rolling.py @@ -124,9 +124,7 @@ def __init__( closed: str | None = None, method: str = "single", ): - # error: Incompatible types in assignment (expression has type "FrameOrSeries", - # variable has type "Union[DataFrame, Series]") - self.obj = obj # type: ignore[assignment] + self.obj = obj self.on = on self.closed = closed self.window = window