TST: series/indexing tests parametrization + moving test methods (#20059)

Author: almaleksia

pandas/tests/series/indexing/test_alter_index.py

@@ -21,88 +21,72 @@
 JOIN_TYPES = ['inner', 'outer', 'left', 'right']
 
 
-def test_align(test_data):
-    def _check_align(a, b, how='left', fill=None):
-        aa, ab = a.align(b, join=how, fill_value=fill)
-
-        join_index = a.index.join(b.index, how=how)
-        if fill is not None:
-            diff_a = aa.index.difference(join_index)
-            diff_b = ab.index.difference(join_index)
-            if len(diff_a) > 0:
-                assert (aa.reindex(diff_a) == fill).all()
-            if len(diff_b) > 0:
-                assert (ab.reindex(diff_b) == fill).all()
-
-        ea = a.reindex(join_index)
-        eb = b.reindex(join_index)
-
-        if fill is not None:
-            ea = ea.fillna(fill)
-            eb = eb.fillna(fill)
-
-        assert_series_equal(aa, ea)
-        assert_series_equal(ab, eb)
-        assert aa.name == 'ts'
-        assert ea.name == 'ts'
-        assert ab.name == 'ts'
-        assert eb.name == 'ts'
-
-    for kind in JOIN_TYPES:
-        _check_align(test_data.ts[2:], test_data.ts[:-5], how=kind)
-        _check_align(test_data.ts[2:], test_data.ts[:-5], how=kind, fill=-1)
-
-        # empty left
-        _check_align(test_data.ts[:0], test_data.ts[:-5], how=kind)
-        _check_align(test_data.ts[:0], test_data.ts[:-5], how=kind, fill=-1)
-
-        # empty right
-        _check_align(test_data.ts[:-5], test_data.ts[:0], how=kind)
-        _check_align(test_data.ts[:-5], test_data.ts[:0], how=kind, fill=-1)
-
-        # both empty
-        _check_align(test_data.ts[:0], test_data.ts[:0], how=kind)
-        _check_align(test_data.ts[:0], test_data.ts[:0], how=kind, fill=-1)
-
-
-def test_align_fill_method(test_data):
-    def _check_align(a, b, how='left', method='pad', limit=None):
-        aa, ab = a.align(b, join=how, method=method, limit=limit)
-
-        join_index = a.index.join(b.index, how=how)
-        ea = a.reindex(join_index)
-        eb = b.reindex(join_index)
-
-        ea = ea.fillna(method=method, limit=limit)
-        eb = eb.fillna(method=method, limit=limit)
-
-        assert_series_equal(aa, ea)
-        assert_series_equal(ab, eb)
-
-    for kind in JOIN_TYPES:
-        for meth in ['pad', 'bfill']:
-            _check_align(test_data.ts[2:], test_data.ts[:-5],
-                         how=kind, method=meth)
-            _check_align(test_data.ts[2:], test_data.ts[:-5],
-                         how=kind, method=meth, limit=1)
-
-            # empty left
-            _check_align(test_data.ts[:0], test_data.ts[:-5],
-                         how=kind, method=meth)
-            _check_align(test_data.ts[:0], test_data.ts[:-5],
-                         how=kind, method=meth, limit=1)
-
-            # empty right
-            _check_align(test_data.ts[:-5], test_data.ts[:0],
-                         how=kind, method=meth)
-            _check_align(test_data.ts[:-5], test_data.ts[:0],
-                         how=kind, method=meth, limit=1)
-
-            # both empty
-            _check_align(test_data.ts[:0], test_data.ts[:0],
-                         how=kind, method=meth)
-            _check_align(test_data.ts[:0], test_data.ts[:0],
-                         how=kind, method=meth, limit=1)
+@pytest.mark.parametrize(
+    'first_slice,second_slice', [
+        [[2, None], [None, -5]],
+        [[None, 0], [None, -5]],
+        [[None, -5], [None, 0]],
+        [[None, 0], [None, 0]]
+    ])
+@pytest.mark.parametrize('join_type', JOIN_TYPES)
+@pytest.mark.parametrize('fill', [None, -1])
+def test_align(test_data, first_slice, second_slice, join_type, fill):
+    a = test_data.ts[slice(*first_slice)]
+    b = test_data.ts[slice(*second_slice)]
+
+    aa, ab = a.align(b, join=join_type, fill_value=fill)
+
+    join_index = a.index.join(b.index, how=join_type)
+    if fill is not None:
+        diff_a = aa.index.difference(join_index)
+        diff_b = ab.index.difference(join_index)
+        if len(diff_a) > 0:
+            assert (aa.reindex(diff_a) == fill).all()
+        if len(diff_b) > 0:
+            assert (ab.reindex(diff_b) == fill).all()
+
+    ea = a.reindex(join_index)
+    eb = b.reindex(join_index)
+
+    if fill is not None:
+        ea = ea.fillna(fill)
+        eb = eb.fillna(fill)
+
+    assert_series_equal(aa, ea)
+    assert_series_equal(ab, eb)
+    assert aa.name == 'ts'
+    assert ea.name == 'ts'
+    assert ab.name == 'ts'
+    assert eb.name == 'ts'
+
+
+@pytest.mark.parametrize(
+    'first_slice,second_slice', [
+        [[2, None], [None, -5]],
+        [[None, 0], [None, -5]],
+        [[None, -5], [None, 0]],
+        [[None, 0], [None, 0]]
+    ])
+@pytest.mark.parametrize('join_type', JOIN_TYPES)
+@pytest.mark.parametrize('method', ['pad', 'bfill'])
+@pytest.mark.parametrize('limit', [None, 1])
+def test_align_fill_method(test_data,
+                           first_slice, second_slice,
+                           join_type, method, limit):
+    a = test_data.ts[slice(*first_slice)]
+    b = test_data.ts[slice(*second_slice)]
+
+    aa, ab = a.align(b, join=join_type, method=method, limit=limit)
+
+    join_index = a.index.join(b.index, how=join_type)
+    ea = a.reindex(join_index)
+    eb = b.reindex(join_index)
+
+    ea = ea.fillna(method=method, limit=limit)
+    eb = eb.fillna(method=method, limit=limit)
+
+    assert_series_equal(aa, ea)
+    assert_series_equal(ab, eb)
 
 
 def test_align_nocopy(test_data):
@@ -481,3 +465,56 @@ def test_rename():
     assert_series_equal(result, expected)
 
     assert result.name == expected.name
+
+
+def test_drop():
+    # unique
+    s = Series([1, 2], index=['one', 'two'])
+    expected = Series([1], index=['one'])
+    result = s.drop(['two'])
+    assert_series_equal(result, expected)
+    result = s.drop('two', axis='rows')
+    assert_series_equal(result, expected)
+
+    # non-unique
+    # GH 5248
+    s = Series([1, 1, 2], index=['one', 'two', 'one'])
+    expected = Series([1, 2], index=['one', 'one'])
+    result = s.drop(['two'], axis=0)
+    assert_series_equal(result, expected)
+    result = s.drop('two')
+    assert_series_equal(result, expected)
+
+    expected = Series([1], index=['two'])
+    result = s.drop(['one'])
+    assert_series_equal(result, expected)
+    result = s.drop('one')
+    assert_series_equal(result, expected)
+
+    # single string/tuple-like
+    s = Series(range(3), index=list('abc'))
+    pytest.raises(KeyError, s.drop, 'bc')
+    pytest.raises(KeyError, s.drop, ('a',))
+
+    # errors='ignore'
+    s = Series(range(3), index=list('abc'))
+    result = s.drop('bc', errors='ignore')
+    assert_series_equal(result, s)
+    result = s.drop(['a', 'd'], errors='ignore')
+    expected = s.iloc[1:]
+    assert_series_equal(result, expected)
+
+    # bad axis
+    pytest.raises(ValueError, s.drop, 'one', axis='columns')
+
+    # GH 8522
+    s = Series([2, 3], index=[True, False])
+    assert s.index.is_object()
+    result = s.drop(True)
+    expected = Series([3], index=[False])
+    assert_series_equal(result, expected)
+
+    # GH 16877
+    s = Series([2, 3], index=[0, 1])
+    with tm.assert_raises_regex(KeyError, 'not contained in axis'):
+        s.drop([False, True])

pandas/tests/series/indexing/test_boolean.py

@@ -283,34 +283,30 @@ def test_where_error():
                   [])
 
 
-def test_where_array_like():
+@pytest.mark.parametrize('klass', [list, tuple, np.array, Series])
+def test_where_array_like(klass):
     # see gh-15414
     s = Series([1, 2, 3])
     cond = [False, True, True]
     expected = Series([np.nan, 2, 3])
-    klasses = [list, tuple, np.array, Series]
 
-    for klass in klasses:
-        result = s.where(klass(cond))
-        assert_series_equal(result, expected)
+    result = s.where(klass(cond))
+    assert_series_equal(result, expected)
 
 
-def test_where_invalid_input():
+@pytest.mark.parametrize('cond', [
+    [1, 0, 1],
+    Series([2, 5, 7]),
+    ["True", "False", "True"],
+    [Timestamp("2017-01-01"), pd.NaT, Timestamp("2017-01-02")]
+])
+def test_where_invalid_input(cond):
     # see gh-15414: only boolean arrays accepted
     s = Series([1, 2, 3])
     msg = "Boolean array expected for the condition"
 
-    conds = [
-        [1, 0, 1],
-        Series([2, 5, 7]),
-        ["True", "False", "True"],
-        [Timestamp("2017-01-01"),
-         pd.NaT, Timestamp("2017-01-02")]
-    ]
-
-    for cond in conds:
-        with tm.assert_raises_regex(ValueError, msg):
-            s.where(cond)
+    with tm.assert_raises_regex(ValueError, msg):
+        s.where(cond)
 
     msg = "Array conditional must be same shape as self"
     with tm.assert_raises_regex(ValueError, msg):
@@ -403,37 +399,43 @@ def f():
     assert_series_equal(s, expected)
 
 
-def test_where_broadcast():
-    # Test a variety of differently sized series
-    for size in range(2, 6):
-        # Test a variety of boolean indices
-        for selection in [
-            # First element should be set
-            np.resize([True, False, False, False, False], size),
-            # Set alternating elements]
-            np.resize([True, False], size),
-            # No element should be set
-            np.resize([False], size)
-        ]:
-
-            # Test a variety of different numbers as content
-            for item in [2.0, np.nan, np.finfo(np.float).max,
-                         np.finfo(np.float).min]:
-                # Test numpy arrays, lists and tuples as the input to be
-                # broadcast
-                for arr in [np.array([item]), [item], (item,)]:
-                    data = np.arange(size, dtype=float)
-                    s = Series(data)
-                    s[selection] = arr
-                    # Construct the expected series by taking the source
-                    # data or item based on the selection
-                    expected = Series([item if use_item else data[
-                        i] for i, use_item in enumerate(selection)])
-                    assert_series_equal(s, expected)
-
-                    s = Series(data)
-                    result = s.where(~selection, arr)
-                    assert_series_equal(result, expected)
+@pytest.mark.parametrize('size', range(2, 6))
+@pytest.mark.parametrize('mask', [
+    [True, False, False, False, False],
+    [True, False],
+    [False]
+])
+@pytest.mark.parametrize('item', [
+    2.0, np.nan, np.finfo(np.float).max, np.finfo(np.float).min
+])
+# Test numpy arrays, lists and tuples as the input to be
+# broadcast
+@pytest.mark.parametrize('box', [
+    lambda x: np.array([x]),
+    lambda x: [x],
+    lambda x: (x,)
+])
+def test_broadcast(size, mask, item, box):
+    selection = np.resize(mask, size)
+
+    data = np.arange(size, dtype=float)
+
+    # Construct the expected series by taking the source
+    # data or item based on the selection
+    expected = Series([item if use_item else data[
+        i] for i, use_item in enumerate(selection)])
+
+    s = Series(data)
+    s[selection] = box(item)
+    assert_series_equal(s, expected)
+
+    s = Series(data)
+    result = s.where(~selection, box(item))
+    assert_series_equal(result, expected)
+
+    s = Series(data)
+    result = s.mask(selection, box(item))
+    assert_series_equal(result, expected)
 
 
 def test_where_inplace():
@@ -587,29 +589,6 @@ def test_mask():
     assert_series_equal(result, expected)
 
 
-def test_mask_broadcast():
-    # GH 8801
-    # copied from test_where_broadcast
-    for size in range(2, 6):
-        for selection in [
-            # First element should be set
-            np.resize([True, False, False, False, False], size),
-            # Set alternating elements]
-            np.resize([True, False], size),
-            # No element should be set
-            np.resize([False], size)
-        ]:
-            for item in [2.0, np.nan, np.finfo(np.float).max,
-                         np.finfo(np.float).min]:
-                for arr in [np.array([item]), [item], (item,)]:
-                    data = np.arange(size, dtype=float)
-                    s = Series(data)
-                    result = s.mask(selection, arr)
-                    expected = Series([item if use_item else data[
-                        i] for i, use_item in enumerate(selection)])
-                    assert_series_equal(result, expected)
-
-
 def test_mask_inplace():
     s = Series(np.random.randn(5))
     cond = s > 0

pandas/tests/series/indexing/test_callable.py

@@ -0,0 +1,33 @@
+import pandas as pd
+import pandas.util.testing as tm
+
+
+def test_getitem_callable():
+    # GH 12533
+    s = pd.Series(4, index=list('ABCD'))
+    result = s[lambda x: 'A']
+    assert result == s.loc['A']
+
+    result = s[lambda x: ['A', 'B']]
+    tm.assert_series_equal(result, s.loc[['A', 'B']])
+
+    result = s[lambda x: [True, False, True, True]]
+    tm.assert_series_equal(result, s.iloc[[0, 2, 3]])
+
+
+def test_setitem_callable():
+    # GH 12533
+    s = pd.Series([1, 2, 3, 4], index=list('ABCD'))
+    s[lambda x: 'A'] = -1
+    tm.assert_series_equal(s, pd.Series([-1, 2, 3, 4], index=list('ABCD')))
+
+
+def test_setitem_other_callable():
+    # GH 13299
+    inc = lambda x: x + 1
+
+    s = pd.Series([1, 2, -1, 4])
+    s[s < 0] = inc
+
+    expected = pd.Series([1, 2, inc, 4])
+    tm.assert_series_equal(s, expected)