diff --git a/pandas/core/frame.py b/pandas/core/frame.py index 6559fc4c24ce2..d2491f3dcceba 100644 --- a/pandas/core/frame.py +++ b/pandas/core/frame.py @@ -2214,7 +2214,7 @@ def query(self, expr, inplace=False, **kwargs): try: new_data = self.loc[res] - except ValueError: + except (ValueError, NotImplementedError): # when res is multi-dimensional loc raises, but this is sometimes a # valid query new_data = self[res] diff --git a/pandas/core/indexing.py b/pandas/core/indexing.py index ae0aaf98fdf02..eb960e70d7987 100755 --- a/pandas/core/indexing.py +++ b/pandas/core/indexing.py @@ -1511,37 +1511,28 @@ def _getitem_axis(self, key, axis=0): elif is_bool_indexer(key): return self._getbool_axis(key, axis=axis) elif is_list_like_indexer(key): + if isinstance(key, ABCDataFrame): + # GH 15438 + raise NotImplementedError("Indexing a with a DataFrame key is " + "not implemented") + elif hasattr(key, 'ndim') and key.ndim > 1: + raise NotImplementedError("Indexing with a multidimensional " + "key is not implemented") # convert various list-like indexers # to a list of keys # we will use the *values* of the object # and NOT the index if its a PandasObject if isinstance(labels, MultiIndex): - - if isinstance(key, (ABCSeries, np.ndarray)) and key.ndim <= 1: - # Series, or 0,1 ndim ndarray + if isinstance(key, (ABCSeries, np.ndarray)) and key.ndim != 1: + # Series or 1-dim ndarray # GH 14730 key = list(key) - elif isinstance(key, ABCDataFrame): - # GH 15438 - raise NotImplementedError("Indexing a MultiIndex with a " - "DataFrame key is not " - "implemented") - elif hasattr(key, 'ndim') and key.ndim > 1: - raise NotImplementedError("Indexing a MultiIndex with a " - "multidimensional key is not " - "implemented") - - if (not isinstance(key, tuple) and len(key) > 1 and - not isinstance(key[0], tuple)): - key = tuple([key]) + if not isinstance(key, tuple): + return self._getitem_iterable(key, axis=axis) # an iterable multi-selection - if not (isinstance(key, tuple) and isinstance(labels, MultiIndex)): - - if hasattr(key, 'ndim') and key.ndim > 1: - raise ValueError('Cannot index with multidimensional key') - + else: return self._getitem_iterable(key, axis=axis) # nested tuple slicing diff --git a/pandas/tests/indexing/test_multiindex.py b/pandas/tests/indexing/test_multiindex.py index c12bb8910ffc9..bedc1e1669ed0 100644 --- a/pandas/tests/indexing/test_multiindex.py +++ b/pandas/tests/indexing/test_multiindex.py @@ -174,6 +174,10 @@ def test_loc_getitem_series(self): result = x.loc[empty] tm.assert_series_equal(result, expected) + with tm.assertRaises(KeyError): + # GH15452 + x.loc[[4, 5]] + def test_loc_getitem_array(self): # GH15434 # passing an array as a key with a MultiIndex @@ -203,6 +207,27 @@ def test_loc_getitem_array(self): result = x.loc[scalar] tm.assert_series_equal(result, expected) + def test_loc_generator(self): + index = MultiIndex.from_product([[1, 2, 3], ['A', 'B', 'C']]) + x = Series(index=index, data=range(9), dtype=np.float64) + y = [1, 3] + + # getitem: + expected = Series( + data=[0, 1, 2, 6, 7, 8], + index=MultiIndex.from_product([[1, 3], ['A', 'B', 'C']]), + dtype=np.float64) + result = x.loc[iter(y)] + tm.assert_series_equal(result, expected) + + # setitem: + expected = Series( + data=[9, 10, 11, 3, 4, 5, 12, 13, 14], + index=index, + dtype=np.float64) + x.loc[iter(y)] = range(9, 15) + tm.assert_series_equal(x, expected) + def test_iloc_getitem_multiindex(self): mi_labels = DataFrame(np.random.randn(4, 3), columns=[['i', 'i', 'j'], ['A', 'A', 'B']], diff --git a/pandas/tests/series/test_indexing.py b/pandas/tests/series/test_indexing.py index 6d8a54b538237..10d06220b4215 100644 --- a/pandas/tests/series/test_indexing.py +++ b/pandas/tests/series/test_indexing.py @@ -4,6 +4,7 @@ import pytest from datetime import datetime, timedelta +import pytest from numpy import nan import numpy as np @@ -290,6 +291,44 @@ def test_getitem_generator(self): assert_series_equal(result, expected) assert_series_equal(result2, expected) + def test_setitem_generator(self): + bool_idx = self.series > 0 + idces = self.series[bool_idx].index + + values = range(bool_idx.sum()) + + expected = self.series.copy() + expected[bool_idx] = values + + # list of labels: + s1 = self.series.copy() + s1[iter(idces)] = values + assert_series_equal(s1, expected) + + # list of labels with .loc: + s2 = self.series.copy() + s2.loc[iter(idces)] = values + assert_series_equal(s2, expected) + + @pytest.mark.xfail(reason="Setitem with booleans generators unsupported") + def test_setitem_boolean_generator(self): + bool_idx = self.series > 0 + + values = range(bool_idx.sum()) + + expected = self.series.copy() + expected[bool_idx] = values + + # boolean generator (fails) + s1 = self.series.copy() + s1[iter(bool_idx)] = values + assert_series_equal(s1, expected) + + # boolean generator with .loc (fails) + s2 = self.series.copy() + s2.loc[iter(bool_idx)] = values + assert_series_equal(s2, expected) + def test_type_promotion(self): # GH12599 s = pd.Series() diff --git a/pandas/tests/sparse/test_indexing.py b/pandas/tests/sparse/test_indexing.py index 382cff4b9d0ac..bfbdc498f1b5b 100644 --- a/pandas/tests/sparse/test_indexing.py +++ b/pandas/tests/sparse/test_indexing.py @@ -512,13 +512,17 @@ def test_loc(self): tm.assert_sp_series_equal(sparse.loc['B'], orig.loc['B'].to_sparse()) - result = sparse.loc[[1, 3, 4]] - exp = orig.loc[[1, 3, 4]].to_sparse() + with tm.assertRaises(KeyError): + # GH15452 + sparse.loc[['D', 'E', 'F']] + + result = sparse.loc[['A', 'B']] + exp = orig.loc[['A', 'B']].to_sparse() tm.assert_sp_series_equal(result, exp) # exceeds the bounds - result = sparse.loc[[1, 3, 4, 5]] - exp = orig.loc[[1, 3, 4, 5]].to_sparse() + result = sparse.loc[['A', 'B', 'C', 'D']] + exp = orig.loc[['A', 'B', 'C', 'D']].to_sparse() tm.assert_sp_series_equal(result, exp) # single element list (GH 15447)