TST: add arithmetic operators fixture (#21033)

Author: jreback

pandas/conftest.py

@@ -2,6 +2,7 @@
 
 import numpy as np
 import pandas as pd
+from pandas.compat import PY3
 import pandas.util._test_decorators as td
 
 
@@ -77,6 +78,24 @@ def observed(request):
     return request.param
 
 
+_all_arithmetic_operators = ['__add__', '__radd__',
+                             '__sub__', '__rsub__',
+                             '__mul__', '__rmul__',
+                             '__floordiv__', '__rfloordiv__',
+                             '__truediv__', '__rtruediv__',
+                             '__pow__', '__rpow__']
+if not PY3:
+    _all_arithmetic_operators.extend(['__div__', '__rdiv__'])
+
+
+@pytest.fixture(params=_all_arithmetic_operators)
+def all_arithmetic_operators(request):
+    """
+    Fixture for dunder names for common arithmetic operations
+    """
+    return request.param
+
+
 @pytest.fixture(params=[None, 'gzip', 'bz2', 'zip',
                         pytest.param('xz', marks=td.skip_if_no_lzma)])
 def compression(request):

pandas/tests/series/test_operators.py

@@ -827,16 +827,52 @@ def test_sub_datetime64_not_ns(self, box, assert_func):
         res = dt64 - obj
         assert_func(res, -expected)
 
-    def test_operators_datetimelike(self):
-        def run_ops(ops, get_ser, test_ser):
+    def test_operators_datetimelike_invalid(self, all_arithmetic_operators):
+        # these are all TypeEror ops
+        op_str = all_arithmetic_operators
+
+        def check(get_ser, test_ser):
 
             # check that we are getting a TypeError
             # with 'operate' (from core/ops.py) for the ops that are not
             # defined
-            for op_str in ops:
-                op = getattr(get_ser, op_str, None)
-                with tm.assert_raises_regex(TypeError, 'operate|cannot'):
-                    op(test_ser)
+            op = getattr(get_ser, op_str, None)
+            with tm.assert_raises_regex(TypeError, 'operate|cannot'):
+                op(test_ser)
+
+        # ## timedelta64 ###
+        td1 = Series([timedelta(minutes=5, seconds=3)] * 3)
+        td1.iloc[2] = np.nan
+
+        # ## datetime64 ###
+        dt1 = Series([Timestamp('20111230'), Timestamp('20120101'),
+                      Timestamp('20120103')])
+        dt1.iloc[2] = np.nan
+        dt2 = Series([Timestamp('20111231'), Timestamp('20120102'),
+                      Timestamp('20120104')])
+        if op_str not in ['__sub__', '__rsub__']:
+            check(dt1, dt2)
+
+        # ## datetime64 with timetimedelta ###
+        # TODO(jreback) __rsub__ should raise?
+        if op_str not in ['__add__', '__radd__', '__sub__']:
+            check(dt1, td1)
+
+        # 8260, 10763
+        # datetime64 with tz
+        tz = 'US/Eastern'
+        dt1 = Series(date_range('2000-01-01 09:00:00', periods=5,
+                                tz=tz), name='foo')
+        dt2 = dt1.copy()
+        dt2.iloc[2] = np.nan
+        td1 = Series(timedelta_range('1 days 1 min', periods=5, freq='H'))
+        td2 = td1.copy()
+        td2.iloc[1] = np.nan
+
+        if op_str not in ['__add__', '__radd__', '__sub__', '__rsub__']:
+            check(dt2, td2)
+
+    def test_operators_datetimelike(self):
 
         # ## timedelta64 ###
         td1 = Series([timedelta(minutes=5, seconds=3)] * 3)
@@ -848,47 +884,31 @@ def run_ops(ops, get_ser, test_ser):
         dt1.iloc[2] = np.nan
         dt2 = Series([Timestamp('20111231'), Timestamp('20120102'),
                       Timestamp('20120104')])
-        ops = ['__add__', '__mul__', '__floordiv__', '__truediv__', '__div__',
-               '__pow__', '__radd__', '__rmul__', '__rfloordiv__',
-               '__rtruediv__', '__rdiv__', '__rpow__']
-        run_ops(ops, dt1, dt2)
         dt1 - dt2
         dt2 - dt1
 
         # ## datetime64 with timetimedelta ###
-        ops = ['__mul__', '__floordiv__', '__truediv__', '__div__', '__pow__',
-               '__rmul__', '__rfloordiv__', '__rtruediv__', '__rdiv__',
-               '__rpow__']
-        run_ops(ops, dt1, td1)
         dt1 + td1
         td1 + dt1
         dt1 - td1
         # TODO: Decide if this ought to work.
         # td1 - dt1
 
         # ## timetimedelta with datetime64 ###
-        ops = ['__sub__', '__mul__', '__floordiv__', '__truediv__', '__div__',
-               '__pow__', '__rmul__', '__rfloordiv__', '__rtruediv__',
-               '__rdiv__', '__rpow__']
-        run_ops(ops, td1, dt1)
         td1 + dt1
         dt1 + td1
 
-        # 8260, 10763
-        # datetime64 with tz
-        ops = ['__mul__', '__floordiv__', '__truediv__', '__div__', '__pow__',
-               '__rmul__', '__rfloordiv__', '__rtruediv__', '__rdiv__',
-               '__rpow__']
+    def test_operators_datetimelike_with_timezones(self):
 
         tz = 'US/Eastern'
         dt1 = Series(date_range('2000-01-01 09:00:00', periods=5,
                                 tz=tz), name='foo')
         dt2 = dt1.copy()
         dt2.iloc[2] = np.nan
+
         td1 = Series(timedelta_range('1 days 1 min', periods=5, freq='H'))
         td2 = td1.copy()
         td2.iloc[1] = np.nan
-        run_ops(ops, dt1, td1)
 
         result = dt1 + td1[0]
         exp = (dt1.dt.tz_localize(None) + td1[0]).dt.tz_localize(tz)
@@ -1133,25 +1153,23 @@ def test_dt64_series_arith_overflow(self):
         res = dt - ser
         tm.assert_series_equal(res, -expected)
 
+    @pytest.mark.parametrize('op', ['__add__', '__radd__',
+                                    '__sub__', '__rsub__'])
     @pytest.mark.parametrize('tz', [None, 'Asia/Tokyo'])
-    def test_dt64_series_add_intlike(self, tz):
+    def test_dt64_series_add_intlike(self, tz, op):
         # GH#19123
         dti = pd.DatetimeIndex(['2016-01-02', '2016-02-03', 'NaT'], tz=tz)
         ser = Series(dti)
 
         other = Series([20, 30, 40], dtype='uint8')
 
-        pytest.raises(TypeError, ser.__add__, 1)
-        pytest.raises(TypeError, ser.__sub__, 1)
+        pytest.raises(TypeError, getattr(ser, op), 1)
 
-        pytest.raises(TypeError, ser.__add__, other)
-        pytest.raises(TypeError, ser.__sub__, other)
+        pytest.raises(TypeError, getattr(ser, op), other)
 
-        pytest.raises(TypeError, ser.__add__, other.values)
-        pytest.raises(TypeError, ser.__sub__, other.values)
+        pytest.raises(TypeError, getattr(ser, op), other.values)
 
-        pytest.raises(TypeError, ser.__add__, pd.Index(other))
-        pytest.raises(TypeError, ser.__sub__, pd.Index(other))
+        pytest.raises(TypeError, getattr(ser, op), pd.Index(other))
 
 
 class TestSeriesOperators(TestData):