Skip to content

REF: implement LocationIndexer._tupleize_axis_indexer #45378

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.

Sign up for GitHub

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

Merged
merged 2 commits into from
Jan 16, 2022
Merged

REF: implement LocationIndexer._tupleize_axis_indexer #45378

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
fc354d8
REF: LocationIndexer._tupleize_axis_indexer
jbrockmendel Jan 14, 2022
44ca0a5
mypy fixup
jbrockmendel Jan 15, 2022
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
50 changes: 35 additions & 15 deletions pandas/core/indexing.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
TYPE_CHECKING,
Hashable,
Sequence,
final,
)
import warnings

Expand Down Expand Up @@ -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)
Expand All @@ -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)

Expand All @@ -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
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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
Expand All @@ -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.
Expand All @@ -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
Expand All @@ -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:
Expand All @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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

Expand Down