REF/TST: Add more pytest idiom to indexing/multiindex/test_getitem.py (#24409)

Author: simonjayhawkins

pandas/tests/indexing/multiindex/test_getitem.py

@@ -4,8 +4,9 @@
 from pandas.compat import lrange, range, u, zip
 
 import pandas as pd
-from pandas import DataFrame, Index, MultiIndex, Series, date_range
+from pandas import DataFrame, Index, MultiIndex, Series
 import pandas.core.common as com
+from pandas.core.indexing import IndexingError
 from pandas.util import testing as tm
 
 
@@ -27,29 +28,6 @@ def test_series_getitem_multiindex(access_method, level1_value, expected):
     tm.assert_series_equal(result, expected)
 
 
-def test_series_getitem_multiindex_xs():
-    # GH6258
-    dt = list(date_range('20130903', periods=3))
-    idx = MultiIndex.from_product([list('AB'), dt])
-    s = Series([1, 3, 4, 1, 3, 4], index=idx)
-
-    result = s.xs('20130903', level=1)
-    expected = Series([1, 1], index=list('AB'))
-    tm.assert_series_equal(result, expected)
-
-
-def test_series_getitem_multiindex_xs_by_label():
-    # GH5684
-    idx = MultiIndex.from_tuples([('a', 'one'), ('a', 'two'), ('b', 'one'),
-                                  ('b', 'two')])
-    s = Series([1, 2, 3, 4], index=idx)
-    s.index.set_names(['L1', 'L2'], inplace=True)
-    result = s.xs('one', level='L2')
-    expected = Series([1, 3], index=['a', 'b'])
-    expected.index.set_names(['L1'], inplace=True)
-    tm.assert_series_equal(result, expected)
-
-
 @pytest.mark.parametrize('level0_value', ['D', 'A'])
 def test_getitem_duplicates_multiindex(level0_value):
     # GH 5725 the 'A' happens to be a valid Timestamp so the doesn't raise
@@ -145,63 +123,83 @@ def test_getitem_duplicates_multiindex_non_scalar_type_object():
 
 
 def test_getitem_simple(multiindex_dataframe_random_data):
-    frame = multiindex_dataframe_random_data
-    df = frame.T
-
-    col = df['foo', 'one']
-    tm.assert_almost_equal(col.values, df.values[:, 0])
-    msg = r"\('foo', 'four'\)"
-    with pytest.raises(KeyError, match=msg):
-        df[('foo', 'four')]
-    msg = "'foobar'"
-    with pytest.raises(KeyError, match=msg):
-        df['foobar']
+    df = multiindex_dataframe_random_data.T
+    expected = df.values[:, 0]
+    result = df['foo', 'one'].values
+    tm.assert_almost_equal(result, expected)
 
 
-@pytest.mark.filterwarnings("ignore:\\n.ix:DeprecationWarning")
-def test_series_getitem(multiindex_year_month_day_dataframe_random_data):
-    ymd = multiindex_year_month_day_dataframe_random_data
-    s = ymd['A']
-
-    result = s[2000, 3]
+@pytest.mark.parametrize('indexer,msg', [
+    (lambda df: df[('foo', 'four')], r"\('foo', 'four'\)"),
+    (lambda df: df['foobar'], "'foobar'")
+])
+def test_getitem_simple_key_error(
+        multiindex_dataframe_random_data, indexer, msg):
+    df = multiindex_dataframe_random_data.T
+    with pytest.raises(KeyError, match=msg):
+        indexer(df)
 
-    # TODO(wesm): unused?
-    # result2 = s.loc[2000, 3]
 
+@pytest.mark.parametrize('indexer', [
+    lambda s: s[2000, 3],
+    lambda s: s.loc[2000, 3]
+])
+def test_series_getitem(
+        multiindex_year_month_day_dataframe_random_data, indexer):
+    s = multiindex_year_month_day_dataframe_random_data['A']
     expected = s.reindex(s.index[42:65])
     expected.index = expected.index.droplevel(0).droplevel(0)
+
+    result = indexer(s)
     tm.assert_series_equal(result, expected)
 
-    result = s[2000, 3, 10]
-    expected = s[49]
+
+@pytest.mark.parametrize('indexer', [
+    lambda s: s[2000, 3, 10],
+    lambda s: s.loc[2000, 3, 10]
+])
+def test_series_getitem_returns_scalar(
+        multiindex_year_month_day_dataframe_random_data, indexer):
+    s = multiindex_year_month_day_dataframe_random_data['A']
+    expected = s.iloc[49]
+
+    result = indexer(s)
     assert result == expected
 
-    # fancy
+
+@pytest.mark.filterwarnings("ignore:\\n.ix:DeprecationWarning")
+@pytest.mark.parametrize('indexer', [
+    lambda s: s.loc[[(2000, 3, 10), (2000, 3, 13)]],
+    lambda s: s.ix[[(2000, 3, 10), (2000, 3, 13)]]
+])
+def test_series_getitem_fancy(
+        multiindex_year_month_day_dataframe_random_data, indexer):
+    s = multiindex_year_month_day_dataframe_random_data['A']
     expected = s.reindex(s.index[49:51])
-    result = s.loc[[(2000, 3, 10), (2000, 3, 13)]]
-    tm.assert_series_equal(result, expected)
 
-    result = s.ix[[(2000, 3, 10), (2000, 3, 13)]]
+    result = indexer(s)
     tm.assert_series_equal(result, expected)
 
-    # key error
-    msg = "356"
-    with pytest.raises(KeyError, match=msg):
-        s.__getitem__((2000, 3, 4))
-
 
-def test_series_getitem_corner(
-        multiindex_year_month_day_dataframe_random_data):
-    ymd = multiindex_year_month_day_dataframe_random_data
-    s = ymd['A']
+@pytest.mark.parametrize('indexer,error,msg', [
+    (lambda s: s.__getitem__((2000, 3, 4)), KeyError, '356'),
+    (lambda s: s[(2000, 3, 4)], KeyError, '356'),
+    (lambda s: s.loc[(2000, 3, 4)], IndexingError, 'Too many indexers'),
+    (lambda s: s.__getitem__(len(s)), IndexError, 'index out of bounds'),
+    (lambda s: s[len(s)], IndexError, 'index out of bounds'),
+    (lambda s: s.iloc[len(s)], IndexError,
+     'single positional indexer is out-of-bounds')
+])
+def test_series_getitem_indexing_errors(
+        multiindex_year_month_day_dataframe_random_data, indexer, error, msg):
+    s = multiindex_year_month_day_dataframe_random_data['A']
+    with pytest.raises(error, match=msg):
+        indexer(s)
 
-    # don't segfault, GH #495
-    # out of bounds access
-    msg = "index out of bounds"
-    with pytest.raises(IndexError, match=msg):
-        s.__getitem__(len(ymd))
 
-    # generator
+def test_series_getitem_corner_generator(
+        multiindex_year_month_day_dataframe_random_data):
+    s = multiindex_year_month_day_dataframe_random_data['A']
     result = s[(x > 0 for x in s)]
     expected = s[s > 0]
     tm.assert_series_equal(result, expected)

pandas/tests/indexing/multiindex/test_xs.py

@@ -3,7 +3,7 @@
 
 from pandas.compat import lrange, product as cart_product
 
-from pandas import DataFrame, Index, MultiIndex, concat
+from pandas import DataFrame, Index, MultiIndex, Series, concat, date_range
 import pandas.core.common as com
 from pandas.util import testing as tm
 
@@ -211,3 +211,27 @@ def test_xs_level_series_slice_not_implemented(
     msg = r'\(2000, slice\(3, 4, None\)\)'
     with pytest.raises(TypeError, match=msg):
         s[2000, 3:4]
+
+
+def test_series_getitem_multiindex_xs():
+    # GH6258
+    dt = list(date_range('20130903', periods=3))
+    idx = MultiIndex.from_product([list('AB'), dt])
+    s = Series([1, 3, 4, 1, 3, 4], index=idx)
+    expected = Series([1, 1], index=list('AB'))
+
+    result = s.xs('20130903', level=1)
+    tm.assert_series_equal(result, expected)
+
+
+def test_series_getitem_multiindex_xs_by_label():
+    # GH5684
+    idx = MultiIndex.from_tuples([('a', 'one'), ('a', 'two'), ('b', 'one'),
+                                  ('b', 'two')])
+    s = Series([1, 2, 3, 4], index=idx)
+    s.index.set_names(['L1', 'L2'], inplace=True)
+    expected = Series([1, 3], index=['a', 'b'])
+    expected.index.set_names(['L1'], inplace=True)
+
+    result = s.xs('one', level='L2')
+    tm.assert_series_equal(result, expected)