From fc354d8d2cccf0809a186b5baa7c70d451b67bfc Mon Sep 17 00:00:00 2001 From: Brock Date: Fri, 14 Jan 2022 13:25:14 -0800 Subject: [PATCH 1/2] REF: LocationIndexer._tupleize_axis_indexer --- pandas/core/indexing.py | 48 ++++++++++++++++++++++++++++------------- 1 file changed, 33 insertions(+), 15 deletions(-) diff --git a/pandas/core/indexing.py b/pandas/core/indexing.py index 77482cbc88bf5..6d82bbc4ac84e 100644 --- a/pandas/core/indexing.py +++ b/pandas/core/indexing.py @@ -5,6 +5,7 @@ TYPE_CHECKING, Hashable, Sequence, + final, ) import warnings @@ -626,6 +627,7 @@ class _LocationIndexer(NDFrameIndexerBase): _valid_types: str axis = None + @final def __call__(self, axis=None): # we need to return a copy of ourselves new_self = type(self)(self.name, self.obj) @@ -640,6 +642,7 @@ def _get_setitem_indexer(self, key): Convert a potentially-label-based key into a positional indexer. """ if self.name == "loc": + # always holds here bc iloc overrides _get_setitem_indexer self._ensure_listlike_indexer(key) if isinstance(key, tuple): @@ -647,7 +650,7 @@ def _get_setitem_indexer(self, key): check_deprecated_indexers(x) if self.axis is not None: - return self._convert_tuple(key) + key = self._tupleize_axis_indexer(key) ax = self.obj._get_axis(0) @@ -658,6 +661,7 @@ def _get_setitem_indexer(self, key): if isinstance(key, tuple): with suppress(IndexingError): + # suppress "Too many indexers" return self._convert_tuple(key) if isinstance(key, range): @@ -665,6 +669,16 @@ def _get_setitem_indexer(self, key): return self._convert_to_indexer(key, axis=0) + @final + def _tupleize_axis_indexer(self, key) -> tuple: + """ + If we have an axis, adapt the given key to be axis-independent. + """ + new_key = [slice(None)] * self.ndim + new_key[self.axis] = key + return tuple(new_key) + + @final def _ensure_listlike_indexer(self, key, axis=None, value=None): """ Ensure that a list-like of column labels are all present by adding them if @@ -702,6 +716,7 @@ def _ensure_listlike_indexer(self, key, axis=None, value=None): keys, axis=0, consolidate=False, only_slice=True ) + @final def __setitem__(self, key, value): check_deprecated_indexers(key) if isinstance(key, tuple): @@ -737,6 +752,7 @@ def _validate_key(self, key, axis: int): """ raise AbstractMethodError(self) + @final def _expand_ellipsis(self, tup: tuple) -> tuple: """ If a tuple key includes an Ellipsis, replace it with an appropriate @@ -758,6 +774,7 @@ def _expand_ellipsis(self, tup: tuple) -> tuple: # by _validate_key_length return tup + @final def _validate_tuple_indexer(self, key: tuple) -> tuple: """ Check the key for valid keys across my indexer. @@ -774,6 +791,7 @@ def _validate_tuple_indexer(self, key: tuple) -> tuple: ) from err return key + @final def _is_nested_tuple_indexer(self, tup: tuple) -> bool: """ Returns @@ -784,23 +802,18 @@ def _is_nested_tuple_indexer(self, tup: tuple) -> bool: return any(is_nested_tuple(tup, ax) for ax in self.obj.axes) return False - def _convert_tuple(self, key): + @final + def _convert_tuple(self, key: tuple) -> tuple: + # Note: we assume _tupleize_axis_indexer has been called, if necessary. keyidx = [] - if self.axis is not None: - axis = self.obj._get_axis_number(self.axis) - for i in range(self.ndim): - if i == axis: - keyidx.append(self._convert_to_indexer(key, axis=axis)) - else: - keyidx.append(slice(None)) - else: - self._validate_key_length(key) - for i, k in enumerate(key): - idx = self._convert_to_indexer(k, axis=i) - keyidx.append(idx) + self._validate_key_length(key) + for i, k in enumerate(key): + idx = self._convert_to_indexer(k, axis=i) + keyidx.append(idx) return tuple(keyidx) + @final def _validate_key_length(self, key: tuple) -> tuple: if len(key) > self.ndim: if key[0] is Ellipsis: @@ -812,6 +825,7 @@ def _validate_key_length(self, key: tuple) -> tuple: raise IndexingError("Too many indexers") return key + @final def _getitem_tuple_same_dim(self, tup: tuple): """ Index with indexers that should return an object of the same dimension @@ -831,6 +845,7 @@ def _getitem_tuple_same_dim(self, tup: tuple): return retval + @final def _getitem_lowerdim(self, tup: tuple): # we can directly get the axis result since the axis is specified @@ -892,6 +907,7 @@ def _getitem_lowerdim(self, tup: tuple): raise IndexingError("not applicable") + @final def _getitem_nested_tuple(self, tup: tuple): # we have a nested tuple so have at least 1 multi-index level # we should be able to match up the dimensionality here @@ -951,6 +967,7 @@ def _getitem_nested_tuple(self, tup: tuple): def _convert_to_indexer(self, key, axis: int): raise AbstractMethodError(self) + @final def __getitem__(self, key): check_deprecated_indexers(key) if type(key) is tuple: @@ -978,6 +995,7 @@ def _getitem_axis(self, key, axis: int): def _has_valid_setitem_indexer(self, indexer) -> bool: raise AbstractMethodError(self) + @final def _getbool_axis(self, key, axis: int): # caller is responsible for ensuring non-None axis labels = self.obj._get_axis(axis) @@ -1544,7 +1562,7 @@ def _get_setitem_indexer(self, key): key = list(key) if self.axis is not None: - return self._convert_tuple(key) + key = self._tupleize_axis_indexer(key) return key From 44ca0a5e983c8a48d3c3d61926f95ce21bd5deac Mon Sep 17 00:00:00 2001 From: Brock Date: Fri, 14 Jan 2022 18:02:43 -0800 Subject: [PATCH 2/2] mypy fixup --- pandas/core/indexing.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/pandas/core/indexing.py b/pandas/core/indexing.py index 6d82bbc4ac84e..9887acd50c6e6 100644 --- a/pandas/core/indexing.py +++ b/pandas/core/indexing.py @@ -675,7 +675,9 @@ def _tupleize_axis_indexer(self, key) -> tuple: If we have an axis, adapt the given key to be axis-independent. """ new_key = [slice(None)] * self.ndim - new_key[self.axis] = key + # error: Invalid index type "Optional[Any]" for "List[slice]"; expected + # type "SupportsIndex" + new_key[self.axis] = key # type:ignore[index] return tuple(new_key) @final