Skip to content
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

Handle .oindex and .vindex for the PandasMultiIndexingAdapter and PandasIndexingAdapter #37

Closed
Closed
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
98 changes: 84 additions & 14 deletions xarray/core/indexing.py
Original file line number Diff line number Diff line change
Expand Up @@ -1680,13 +1680,29 @@ def _convert_scalar(self, item):
# a NumPy array.
return to_0d_array(item)

def _oindex_get(self, indexer: OuterIndexer):
return self.__getitem__(indexer)
def _prepare_key(self, key: tuple[Any, ...]) -> tuple[Any, ...]:
if isinstance(key, tuple) and len(key) == 1:
# unpack key so it can index a pandas.Index object (pandas.Index
# objects don't like tuples)
(key,) = key

def _vindex_get(self, indexer: VectorizedIndexer):
return self.__getitem__(indexer)
return key

def __getitem__(
def _handle_result(
self, result: Any
) -> (
PandasIndexingAdapter
| NumpyIndexingAdapter
| np.ndarray
| np.datetime64
| np.timedelta64
):
if isinstance(result, pd.Index):
return type(self)(result, dtype=self.dtype)
else:
return self._convert_scalar(result)

def _get_item(
self, indexer: ExplicitIndexer
) -> (
PandasIndexingAdapter
Expand All @@ -1695,22 +1711,48 @@ def __getitem__(
| np.datetime64
| np.timedelta64
):
key = indexer.tuple
if isinstance(key, tuple) and len(key) == 1:
# unpack key so it can index a pandas.Index object (pandas.Index
# objects don't like tuples)
(key,) = key
key = self._prepare_key(indexer.tuple)

if getattr(key, "ndim", 0) > 1: # Return np-array if multidimensional
indexable = NumpyIndexingAdapter(np.asarray(self))
return apply_indexer(indexable, indexer)

result = self.array[key]

if isinstance(result, pd.Index):
return type(self)(result, dtype=self.dtype)
else:
return self._convert_scalar(result)
return self._handle_result(result)

def _oindex_get(
self, indexer: OuterIndexer
) -> (
PandasIndexingAdapter
| NumpyIndexingAdapter
| np.ndarray
| np.datetime64
| np.timedelta64
):
return self._get_item(indexer)

def _vindex_get(
self, indexer: VectorizedIndexer
) -> (
PandasIndexingAdapter
| NumpyIndexingAdapter
| np.ndarray
| np.datetime64
| np.timedelta64
):
return self._get_item(indexer)

def __getitem__(
self, indexer: ExplicitIndexer
) -> (
PandasIndexingAdapter
| NumpyIndexingAdapter
| np.ndarray
| np.datetime64
| np.timedelta64
):
return self._get_item(indexer)

def transpose(self, order) -> pd.Index:
return self.array # self.array should be always one-dimensional
Expand Down Expand Up @@ -1766,6 +1808,34 @@ def _convert_scalar(self, item):
item = item[idx]
return super()._convert_scalar(item)

def _oindex_get(
self, indexer: OuterIndexer
) -> (
PandasIndexingAdapter
| NumpyIndexingAdapter
| np.ndarray
| np.datetime64
| np.timedelta64
):
result = super()._oindex_get(indexer)
if isinstance(result, type(self)):
result.level = self.level
return result

def _vindex_get(
self, indexer: VectorizedIndexer
) -> (
PandasIndexingAdapter
| NumpyIndexingAdapter
| np.ndarray
| np.datetime64
| np.timedelta64
):
result = super()._vindex_get(indexer)
if isinstance(result, type(self)):
result.level = self.level
return result

def __getitem__(self, indexer: ExplicitIndexer):
result = super().__getitem__(indexer)
if isinstance(result, type(self)):
Expand Down