From 813496f45eff29375a5b82f00eeffb2fbee87627 Mon Sep 17 00:00:00 2001 From: jbrockmendel Date: Wed, 5 Feb 2020 18:04:03 -0800 Subject: [PATCH 1/3] REF: turn _try_mi into _get_values_for_loc --- pandas/core/indexes/base.py | 6 ++++-- pandas/core/indexes/extension.py | 2 +- pandas/core/indexes/multi.py | 34 ++++++++++++++++++-------------- pandas/core/indexes/numeric.py | 2 +- 4 files changed, 25 insertions(+), 19 deletions(-) diff --git a/pandas/core/indexes/base.py b/pandas/core/indexes/base.py index 891ae95db65a0..56fd097ee9a5d 100644 --- a/pandas/core/indexes/base.py +++ b/pandas/core/indexes/base.py @@ -4601,14 +4601,16 @@ def get_value(self, series: "Series", key): else: raise - return self._get_values_for_loc(series, loc) + return self._get_values_for_loc(series, loc, key) - def _get_values_for_loc(self, series: "Series", loc): + def _get_values_for_loc(self, series: "Series", loc, key): """ Do a positional lookup on the given Series, returning either a scalar or a Series. Assumes that `series.index is self` + + key is included for MultiIndex compat. """ if is_integer(loc): return series._values[loc] diff --git a/pandas/core/indexes/extension.py b/pandas/core/indexes/extension.py index c32889a9360bc..f2338da43dfc0 100644 --- a/pandas/core/indexes/extension.py +++ b/pandas/core/indexes/extension.py @@ -317,6 +317,6 @@ def get_value(self, series: "Series", key): else: raise - return self._get_values_for_loc(series, loc) + return self._get_values_for_loc(series, loc, key) # -------------------------------------------------------------------- diff --git a/pandas/core/indexes/multi.py b/pandas/core/indexes/multi.py index 708bea7d132a2..5ee319c522267 100644 --- a/pandas/core/indexes/multi.py +++ b/pandas/core/indexes/multi.py @@ -2326,27 +2326,31 @@ def get_value(self, series, key): # We have to explicitly exclude generators, as these are hashable. raise InvalidIndexError(key) - def _try_mi(k): - # TODO: what if a level contains tuples?? - loc = self.get_loc(k) - - new_values = series._values[loc] - if is_scalar(loc): - return new_values - - new_index = self[loc] - new_index = maybe_droplevels(new_index, k) - return series._constructor( - new_values, index=new_index, name=series.name - ).__finalize__(self) - try: - return _try_mi(key) + loc = self.get_loc(key) except KeyError: if is_integer(key): return series._values[key] else: raise + else: + return self._get_values_for_loc(series, loc, key) + + def _get_values_for_loc(self, series: "Series", loc, key): + """ + Do a positional lookup on the given Series, returning either a scalar + or a Series. + + Assumes that `series.index is self` + """ + new_values = series._values[loc] + if is_scalar(loc): + return new_values + + new_index = self[loc] + new_index = maybe_droplevels(new_index, key) + new_ser = series._constructor(new_values, index=new_index, name=series.name) + return new_ser.__finalize__(series) def _convert_listlike_indexer(self, keyarr): """ diff --git a/pandas/core/indexes/numeric.py b/pandas/core/indexes/numeric.py index ebfe50327b479..7cc1e6593c2c0 100644 --- a/pandas/core/indexes/numeric.py +++ b/pandas/core/indexes/numeric.py @@ -425,7 +425,7 @@ def get_value(self, series: "Series", key): raise InvalidIndexError loc = self.get_loc(key) - return self._get_values_for_loc(series, loc) + return self._get_values_for_loc(series, loc, key) def equals(self, other) -> bool: """ From 308fe271a01024109ba66047ae806533e01c6009 Mon Sep 17 00:00:00 2001 From: jbrockmendel Date: Wed, 5 Feb 2020 20:26:37 -0800 Subject: [PATCH 2/3] make MultiIndex.get_value more like base class --- pandas/core/indexes/multi.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pandas/core/indexes/multi.py b/pandas/core/indexes/multi.py index 5ee319c522267..1fba6cdacf81a 100644 --- a/pandas/core/indexes/multi.py +++ b/pandas/core/indexes/multi.py @@ -2330,11 +2330,11 @@ def get_value(self, series, key): loc = self.get_loc(key) except KeyError: if is_integer(key): - return series._values[key] + loc = key else: raise - else: - return self._get_values_for_loc(series, loc, key) + + return self._get_values_for_loc(series, loc, key) def _get_values_for_loc(self, series: "Series", loc, key): """ From 3b32c5ae770e990963938bca7bd7677da01e3b4f Mon Sep 17 00:00:00 2001 From: jbrockmendel Date: Wed, 5 Feb 2020 21:39:15 -0800 Subject: [PATCH 3/3] lint fixup --- pandas/core/indexes/multi.py | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/pandas/core/indexes/multi.py b/pandas/core/indexes/multi.py index 1fba6cdacf81a..81e89441e92d8 100644 --- a/pandas/core/indexes/multi.py +++ b/pandas/core/indexes/multi.py @@ -1,5 +1,15 @@ from sys import getsizeof -from typing import Any, Hashable, Iterable, List, Optional, Sequence, Tuple, Union +from typing import ( + TYPE_CHECKING, + Any, + Hashable, + Iterable, + List, + Optional, + Sequence, + Tuple, + Union, +) import warnings import numpy as np @@ -56,6 +66,9 @@ pprint_thing, ) +if TYPE_CHECKING: + from pandas import Series # noqa:F401 + _index_doc_kwargs = dict(ibase._index_doc_kwargs) _index_doc_kwargs.update( dict(klass="MultiIndex", target_klass="MultiIndex or list of tuples")