From ab42f85f192ab054a18f94825ced1bb4c1ab7d3f Mon Sep 17 00:00:00 2001 From: jbrockmendel Date: Sun, 16 Jan 2022 12:24:59 -0800 Subject: [PATCH] REF: implement LocationIndexer._tupleize_axis_indexer (#45378) --- pandas/core/indexing.py | 50 ++++++++++++++++++++++++++++------------- 1 file changed, 35 insertions(+), 15 deletions(-) diff --git a/pandas/core/indexing.py b/pandas/core/indexing.py index 0bf6f22c2d385..30f21dde7f0be 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,18 @@ 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 + # 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 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 +718,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 +754,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 +776,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 +793,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 +804,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 +827,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 +847,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 +909,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 +969,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 +997,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 +1564,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