REF: implement LocationIndexer._tupleize_axis_indexer (#45378)

Author: jbrockmendel

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,14 +642,15 @@ 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):
             for x in 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,13 +661,26 @@ 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):
             return list(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