diff --git a/pandas/core/base.py b/pandas/core/base.py index f30430dd394ca..3911cbb63c07b 100644 --- a/pandas/core/base.py +++ b/pandas/core/base.py @@ -10,7 +10,9 @@ Callable, Dict, FrozenSet, + Hashable, Optional, + Set, TypeVar, Union, cast, @@ -22,6 +24,7 @@ from pandas._typing import ( Dtype, DtypeObj, + FrameOrSeriesUnion, IndexLabel, Shape, ) @@ -167,10 +170,14 @@ class SpecificationError(Exception): class SelectionMixin: """ - mixin implementing the selection & aggregation interface on a group-like - object sub-classes need to define: obj, exclusions + Mixin for the selection & aggregation interface on a group-like object. + + Sub-classes need to define: obj, exclusions """ + obj: FrameOrSeriesUnion + exclusions: Set[Hashable] + _selection: Optional[IndexLabel] = None _internal_names = ["_cache", "__setstate__"] _internal_names_set = set(_internal_names) @@ -225,15 +232,9 @@ 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] - else: - # error: "SelectionMixin" has no attribute "obj" - return self.obj[self._selection] # type: ignore[attr-defined] + if self._selection is None or isinstance(self.obj, ABCSeries): + return self.obj + return self.obj[self._selection] @cache_readonly def ndim(self) -> int: @@ -241,49 +242,31 @@ def ndim(self) -> int: @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/window/rolling.py b/pandas/core/window/rolling.py index 6db86b940737e..4254a8dee12a9 100644 --- a/pandas/core/window/rolling.py +++ b/pandas/core/window/rolling.py @@ -16,7 +16,6 @@ Dict, List, Optional, - Set, Tuple, Type, Union, @@ -115,7 +114,6 @@ class BaseWindow(SelectionMixin): """Provides utilities for performing windowing operations.""" _attributes: List[str] = [] - exclusions: Set[str] = set() def __init__( self, @@ -130,6 +128,7 @@ def __init__( method: str = "single", ): self.obj = obj + self.exclusions = set() self.on = on self.closed = closed self.window = window