parametrize the first few arithmetic tests (#18133)

Author: jbrockmendel

pandas/tests/indexes/datetimes/test_arithmetic.py

@@ -14,77 +14,118 @@
                     date_range)
 
 
+@pytest.fixture(params=[None, 'UTC', 'Asia/Tokyo',
+                        'US/Eastern', 'dateutil/Asia/Singapore',
+                        'dateutil/US/Pacific'])
+def tz(request):
+    return request.param
+
+
+@pytest.fixture(params=[pd.offsets.Hour(2), timedelta(hours=2),
+                        np.timedelta64(2, 'h'), Timedelta(hours=2)],
+                ids=str)
+def delta(request):
+    # Several ways of representing two hours
+    return request.param
+
+
+@pytest.fixture(
+    params=[
+        datetime(2011, 1, 1),
+        DatetimeIndex(['2011-01-01', '2011-01-02']),
+        DatetimeIndex(['2011-01-01', '2011-01-02']).tz_localize('US/Eastern'),
+        np.datetime64('2011-01-01'),
+        Timestamp('2011-01-01')],
+    ids=lambda x: type(x).__name__)
+def addend(request):
+    return request.param
+
+
 class TestDatetimeIndexArithmetic(object):
-    tz = [None, 'UTC', 'Asia/Tokyo', 'US/Eastern', 'dateutil/Asia/Singapore',
-          'dateutil/US/Pacific']
-
-    def test_add_iadd(self):
-        for tz in self.tz:
-
-            # offset
-            offsets = [pd.offsets.Hour(2), timedelta(hours=2),
-                       np.timedelta64(2, 'h'), Timedelta(hours=2)]
-
-            for delta in offsets:
-                rng = pd.date_range('2000-01-01', '2000-02-01', tz=tz)
-                result = rng + delta
-                expected = pd.date_range('2000-01-01 02:00',
-                                         '2000-02-01 02:00', tz=tz)
-                tm.assert_index_equal(result, expected)
-                rng += delta
-                tm.assert_index_equal(rng, expected)
-
-            # int
-            rng = pd.date_range('2000-01-01 09:00', freq='H', periods=10,
-                                tz=tz)
-            result = rng + 1
-            expected = pd.date_range('2000-01-01 10:00', freq='H', periods=10,
-                                     tz=tz)
-            tm.assert_index_equal(result, expected)
-            rng += 1
-            tm.assert_index_equal(rng, expected)
 
+    def test_dti_add_timestamp_raises(self):
         idx = DatetimeIndex(['2011-01-01', '2011-01-02'])
         msg = "cannot add DatetimeIndex and Timestamp"
         with tm.assert_raises_regex(TypeError, msg):
             idx + Timestamp('2011-01-01')
 
+    def test_dti_radd_timestamp_raises(self):
+        idx = DatetimeIndex(['2011-01-01', '2011-01-02'])
+        msg = "cannot add DatetimeIndex and Timestamp"
         with tm.assert_raises_regex(TypeError, msg):
             Timestamp('2011-01-01') + idx
 
-    def test_sub_isub(self):
-        for tz in self.tz:
-
-            # offset
-            offsets = [pd.offsets.Hour(2), timedelta(hours=2),
-                       np.timedelta64(2, 'h'), Timedelta(hours=2)]
-
-            for delta in offsets:
-                rng = pd.date_range('2000-01-01', '2000-02-01', tz=tz)
-                expected = pd.date_range('1999-12-31 22:00',
-                                         '2000-01-31 22:00', tz=tz)
-
-                result = rng - delta
-                tm.assert_index_equal(result, expected)
-                rng -= delta
-                tm.assert_index_equal(rng, expected)
-
-            # int
-            rng = pd.date_range('2000-01-01 09:00', freq='H', periods=10,
-                                tz=tz)
-            result = rng - 1
-            expected = pd.date_range('2000-01-01 08:00', freq='H', periods=10,
-                                     tz=tz)
-            tm.assert_index_equal(result, expected)
-            rng -= 1
-            tm.assert_index_equal(rng, expected)
-
-    @pytest.mark.parametrize('addend', [
-        datetime(2011, 1, 1),
-        DatetimeIndex(['2011-01-01', '2011-01-02']),
-        DatetimeIndex(['2011-01-01', '2011-01-02']).tz_localize('US/Eastern'),
-        np.datetime64('2011-01-01'),
-        Timestamp('2011-01-01')])
+    # -------------------------------------------------------------
+    # Binary operations DatetimeIndex and int
+
+    def test_dti_add_int(self, tz):
+        rng = pd.date_range('2000-01-01 09:00', freq='H',
+                            periods=10, tz=tz)
+        result = rng + 1
+        expected = pd.date_range('2000-01-01 10:00', freq='H',
+                                 periods=10, tz=tz)
+        tm.assert_index_equal(result, expected)
+
+    def test_dti_iadd_int(self, tz):
+        rng = pd.date_range('2000-01-01 09:00', freq='H',
+                            periods=10, tz=tz)
+        expected = pd.date_range('2000-01-01 10:00', freq='H',
+                                 periods=10, tz=tz)
+        rng += 1
+        tm.assert_index_equal(rng, expected)
+
+    def test_dti_sub_int(self, tz):
+        rng = pd.date_range('2000-01-01 09:00', freq='H',
+                            periods=10, tz=tz)
+        result = rng - 1
+        expected = pd.date_range('2000-01-01 08:00', freq='H',
+                                 periods=10, tz=tz)
+        tm.assert_index_equal(result, expected)
+
+    def test_dti_isub_int(self, tz):
+        rng = pd.date_range('2000-01-01 09:00', freq='H',
+                            periods=10, tz=tz)
+        expected = pd.date_range('2000-01-01 08:00', freq='H',
+                                 periods=10, tz=tz)
+        rng -= 1
+        tm.assert_index_equal(rng, expected)
+
+    # -------------------------------------------------------------
+    # Binary operations DatetimeIndex and timedelta-like
+
+    def test_dti_add_timedeltalike(self, tz, delta):
+        rng = pd.date_range('2000-01-01', '2000-02-01', tz=tz)
+        result = rng + delta
+        expected = pd.date_range('2000-01-01 02:00',
+                                 '2000-02-01 02:00', tz=tz)
+        tm.assert_index_equal(result, expected)
+
+    def test_dti_iadd_timedeltalike(self, tz, delta):
+        rng = pd.date_range('2000-01-01', '2000-02-01', tz=tz)
+        expected = pd.date_range('2000-01-01 02:00',
+                                 '2000-02-01 02:00', tz=tz)
+        rng += delta
+        tm.assert_index_equal(rng, expected)
+
+    def test_dti_sub_timedeltalike(self, tz, delta):
+        rng = pd.date_range('2000-01-01', '2000-02-01', tz=tz)
+        expected = pd.date_range('1999-12-31 22:00',
+                                 '2000-01-31 22:00', tz=tz)
+        result = rng - delta
+        tm.assert_index_equal(result, expected)
+
+    def test_dti_isub_timedeltalike(self, tz, delta):
+        rng = pd.date_range('2000-01-01', '2000-02-01', tz=tz)
+        expected = pd.date_range('1999-12-31 22:00',
+                                 '2000-01-31 22:00', tz=tz)
+        rng -= delta
+        tm.assert_index_equal(rng, expected)
+
+    # -------------------------------------------------------------
+    # Binary Operations DatetimeIndex and datetime-like
+    # TODO: A couple other tests belong in this section.  Move them in
+    # A PR where there isn't already a giant diff.
+
     def test_add_datetimelike_and_dti(self, addend):
         # GH#9631
         dti = DatetimeIndex(['2011-01-01', '2011-01-02'])
@@ -95,12 +136,6 @@ def test_add_datetimelike_and_dti(self, addend):
         with tm.assert_raises_regex(TypeError, msg):
             addend + dti
 
-    @pytest.mark.parametrize('addend', [
-        datetime(2011, 1, 1),
-        DatetimeIndex(['2011-01-01', '2011-01-02']),
-        DatetimeIndex(['2011-01-01', '2011-01-02']).tz_localize('US/Eastern'),
-        np.datetime64('2011-01-01'),
-        Timestamp('2011-01-01')])
     def test_add_datetimelike_and_dti_tz(self, addend):
         # GH#9631
         dti_tz = DatetimeIndex(['2011-01-01',
@@ -112,6 +147,8 @@ def test_add_datetimelike_and_dti_tz(self, addend):
         with tm.assert_raises_regex(TypeError, msg):
             addend + dti_tz
 
+    # -------------------------------------------------------------
+
     def test_sub_dti_dti(self):
         # previously performed setop (deprecated in 0.16.0), now changed to
         # return subtraction -> TimeDeltaIndex (GH ...)

pandas/tests/indexes/timedeltas/test_arithmetic.py

@@ -12,9 +12,21 @@
                     Timestamp, Timedelta)
 
 
+@pytest.fixture(params=[pd.offsets.Hour(2), timedelta(hours=2),
+                        np.timedelta64(2, 'h'), Timedelta(hours=2)],
+                ids=str)
+def delta(request):
+    # Several ways of representing two hours
+    return request.param
+
+
+@pytest.fixture(params=['B', 'D'])
+def freq(request):
+    return request.param
+
+
 class TestTimedeltaIndexArithmetic(object):
     _holder = TimedeltaIndex
-    _multiprocess_can_split_ = True
 
     # TODO: Split by ops, better name
     def test_numeric_compat(self):
@@ -88,62 +100,91 @@ def test_ufunc_coercions(self):
             tm.assert_index_equal(result, exp)
             assert result.freq is None
 
-    def test_add_iadd(self):
-        # only test adding/sub offsets as + is now numeric
-
-        # offset
-        offsets = [pd.offsets.Hour(2), timedelta(hours=2),
-                   np.timedelta64(2, 'h'), Timedelta(hours=2)]
-
-        for delta in offsets:
-            rng = timedelta_range('1 days', '10 days')
-            result = rng + delta
-            expected = timedelta_range('1 days 02:00:00', '10 days 02:00:00',
-                                       freq='D')
-            tm.assert_index_equal(result, expected)
-            rng += delta
-            tm.assert_index_equal(rng, expected)
+    # -------------------------------------------------------------
+    # Binary operations TimedeltaIndex and integer
 
-        # int
+    def test_tdi_add_int(self):
         rng = timedelta_range('1 days 09:00:00', freq='H', periods=10)
         result = rng + 1
         expected = timedelta_range('1 days 10:00:00', freq='H', periods=10)
         tm.assert_index_equal(result, expected)
+
+    def test_tdi_iadd_int(self):
+        rng = timedelta_range('1 days 09:00:00', freq='H', periods=10)
+        expected = timedelta_range('1 days 10:00:00', freq='H', periods=10)
         rng += 1
         tm.assert_index_equal(rng, expected)
 
-    def test_sub_isub(self):
-        # only test adding/sub offsets as - is now numeric
-
-        # offset
-        offsets = [pd.offsets.Hour(2), timedelta(hours=2),
-                   np.timedelta64(2, 'h'), Timedelta(hours=2)]
-
-        for delta in offsets:
-            rng = timedelta_range('1 days', '10 days')
-            result = rng - delta
-            expected = timedelta_range('0 days 22:00:00', '9 days 22:00:00')
-            tm.assert_index_equal(result, expected)
-            rng -= delta
-            tm.assert_index_equal(rng, expected)
-
-        # int
+    def test_tdi_sub_int(self):
         rng = timedelta_range('1 days 09:00:00', freq='H', periods=10)
         result = rng - 1
         expected = timedelta_range('1 days 08:00:00', freq='H', periods=10)
         tm.assert_index_equal(result, expected)
+
+    def test_tdi_isub_int(self):
+        rng = timedelta_range('1 days 09:00:00', freq='H', periods=10)
+        expected = timedelta_range('1 days 08:00:00', freq='H', periods=10)
         rng -= 1
         tm.assert_index_equal(rng, expected)
 
+    # -------------------------------------------------------------
+    # Binary operations TimedeltaIndex and timedelta-like
+
+    def test_tdi_add_timedeltalike(self, delta):
+        # only test adding/sub offsets as + is now numeric
+        rng = timedelta_range('1 days', '10 days')
+        result = rng + delta
+        expected = timedelta_range('1 days 02:00:00', '10 days 02:00:00',
+                                   freq='D')
+        tm.assert_index_equal(result, expected)
+
+    def test_tdi_iadd_timedeltalike(self, delta):
+        # only test adding/sub offsets as + is now numeric
+        rng = timedelta_range('1 days', '10 days')
+        expected = timedelta_range('1 days 02:00:00', '10 days 02:00:00',
+                                   freq='D')
+        rng += delta
+        tm.assert_index_equal(rng, expected)
+
+    def test_tdi_sub_timedeltalike(self, delta):
+        # only test adding/sub offsets as - is now numeric
+        rng = timedelta_range('1 days', '10 days')
+        result = rng - delta
+        expected = timedelta_range('0 days 22:00:00', '9 days 22:00:00')
+        tm.assert_index_equal(result, expected)
+
+    def test_tdi_isub_timedeltalike(self, delta):
+        # only test adding/sub offsets as - is now numeric
+        rng = timedelta_range('1 days', '10 days')
+        expected = timedelta_range('0 days 22:00:00', '9 days 22:00:00')
+        rng -= delta
+        tm.assert_index_equal(rng, expected)
+
+    # -------------------------------------------------------------
+    # Binary operations TimedeltaIndex and datetime-like
+
+    def test_tdi_sub_timestamp_raises(self):
         idx = TimedeltaIndex(['1 day', '2 day'])
         msg = "cannot subtract a datelike from a TimedeltaIndex"
         with tm.assert_raises_regex(TypeError, msg):
             idx - Timestamp('2011-01-01')
 
+    def test_tdi_add_timestamp(self):
+        idx = TimedeltaIndex(['1 day', '2 day'])
+
+        result = idx + Timestamp('2011-01-01')
+        expected = DatetimeIndex(['2011-01-02', '2011-01-03'])
+        tm.assert_index_equal(result, expected)
+
+    def test_tdi_radd_timestamp(self):
+        idx = TimedeltaIndex(['1 day', '2 day'])
+
         result = Timestamp('2011-01-01') + idx
         expected = DatetimeIndex(['2011-01-02', '2011-01-03'])
         tm.assert_index_equal(result, expected)
 
+    # -------------------------------------------------------------
+
     # TODO: Split by operation, better name
     def test_ops_compat(self):
 
@@ -634,7 +675,6 @@ def test_tdi_ops_attributes(self):
 
     # TODO: Needs more informative name, probably split up into
     # more targeted tests
-    @pytest.mark.parametrize('freq', ['B', 'D'])
     def test_timedelta(self, freq):
         index = date_range('1/1/2000', periods=50, freq=freq)
 

pandas/tests/test_downstream.py

@@ -1,3 +1,4 @@
+# -*- coding: utf-8 -*-
 """
 Testing that we work in the downstream packages
 """