parametrize the first few arithmetic tests

Author: jbrockmendel

pandas/tests/indexes/datetimes/test_arithmetic.py

@@ -14,70 +14,109 @@
                     date_range)
 
 
+
 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)
+    # Several ways of representing two hours
+    two_hour_variants = [pd.offsets.Hour(2), timedelta(hours=2),
+                         np.timedelta64(2, 'h'), Timedelta(hours=2)]
 
+    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)
+    # -------------------------------------------------------------
+    # Binary operations DatetimeIndex and int
+
+    @pytest.mark.parametrize('tz', tz)
+    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)
+
+    @pytest.mark.parametrize('tz', tz)
+    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)
+
+    @pytest.mark.parametrize('tz', tz)
+    def test_dti_sub_int(self):
+        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)
+
+    @pytest.mark.parametrize('tz', tz)
+    def test_dti_isub_int(self):
+        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
+
+    @pytest.mark.parametrize('delta', two_hour_variants)
+    @pytest.mark.parametrize('tz', tz)
+    def test_dti_add_timedeltalike(self, tz):
+        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)
+
+    @pytest.mark.parametrize('delta', two_hour_variants)
+    @pytest.mark.parametrize('tz', tz)
+    def test_dti_iadd_timedeltalike(self, tz):
+        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)
+
+    @pytest.mark.parametrize('delta', two_hour_variants)
+    @pytest.mark.parametrize('tz', tz)
+    def test_dti_sub_timedeltalike(self, tz):
+        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)
+
+    @pytest.mark.parametrize('delta', two_hour_variants)
+    @pytest.mark.parametrize('tz', tz)
+    def test_dti_isub_timedeltalike(self, tz):
+        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.
 
     @pytest.mark.parametrize('addend', [
         datetime(2011, 1, 1),
@@ -112,6 +151,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

@@ -16,6 +16,10 @@ class TestTimedeltaIndexArithmetic(object):
     _holder = TimedeltaIndex
     _multiprocess_can_split_ = True
 
+    # Several ways of representing two hours
+    two_hour_variants = [pd.offsets.Hour(2), timedelta(hours=2),
+                         np.timedelta64(2, 'h'), Timedelta(hours=2)]
+
     # TODO: Split by ops, better name
     def test_numeric_compat(self):
         idx = self._holder(np.arange(5, dtype='int64'))
@@ -88,62 +92,96 @@ 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
+
+    @pytest.mark.parametrize('delta', two_hour_variants)
+    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)
+
+    @pytest.mark.parametrize('delta', two_hour_variants)
+    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)
+
+    @pytest.mark.parametrize('delta', two_hour_variants)
+    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)
+
+    @pytest.mark.parametrize('delta', two_hour_variants)
+    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):