Skip to content

TYP: handle mypy ignore in SelectionMixin #38239

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 21 commits into from
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
59 changes: 21 additions & 38 deletions pandas/core/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,9 @@
Callable,
Dict,
FrozenSet,
Hashable,
Optional,
Set,
TypeVar,
Union,
cast,
Expand All @@ -22,6 +24,7 @@
from pandas._typing import (
Dtype,
DtypeObj,
FrameOrSeriesUnion,
IndexLabel,
Shape,
)
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -225,65 +232,41 @@ 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:
return self._selected_obj.ndim

@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)

Expand Down
3 changes: 1 addition & 2 deletions pandas/core/window/rolling.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,6 @@
Dict,
List,
Optional,
Set,
Tuple,
Type,
Union,
Expand Down Expand Up @@ -115,7 +114,6 @@ class BaseWindow(SelectionMixin):
"""Provides utilities for performing windowing operations."""

_attributes: List[str] = []
exclusions: Set[str] = set()

def __init__(
self,
Expand All @@ -130,6 +128,7 @@ def __init__(
method: str = "single",
):
self.obj = obj
self.exclusions = set()
self.on = on
self.closed = closed
self.window = window
Expand Down