Refactor _TimeOp._validate to separate datetime vs timedelta vs dateoffset (#18832)

Author: jbrockmendel

pandas/core/ops.py

@@ -387,54 +387,19 @@ def __init__(self, left, right, name, na_op):
         self.lvalues, self.rvalues = self._convert_for_datetime(lvalues,
                                                                 rvalues)
 
-    def _validate(self, lvalues, rvalues, name):
-        # timedelta and integer mul/div
-
-        if ((self.is_timedelta_lhs and
-                (self.is_integer_rhs or self.is_floating_rhs)) or
-            (self.is_timedelta_rhs and
-                (self.is_integer_lhs or self.is_floating_lhs))):
-
-            if name not in ('__div__', '__truediv__', '__mul__', '__rmul__'):
-                raise TypeError("can only operate on a timedelta and an "
-                                "integer or a float for division and "
-                                "multiplication, but the operator [{name}] "
-                                "was passed".format(name=name))
-
-        # 2 timedeltas
-        elif ((self.is_timedelta_lhs and
-               (self.is_timedelta_rhs or self.is_offset_rhs)) or
-              (self.is_timedelta_rhs and
-               (self.is_timedelta_lhs or self.is_offset_lhs))):
-
-            if name not in ('__div__', '__rdiv__', '__truediv__',
-                            '__rtruediv__', '__add__', '__radd__', '__sub__',
-                            '__rsub__'):
-                raise TypeError("can only operate on a timedeltas for addition"
-                                ", subtraction, and division, but the operator"
-                                " [{name}] was passed".format(name=name))
-
-        # datetime and timedelta/DateOffset
-        elif (self.is_datetime_lhs and
-              (self.is_timedelta_rhs or self.is_offset_rhs)):
+    def _validate_datetime(self, lvalues, rvalues, name):
+        # assumes self.is_datetime_lhs
 
+        if (self.is_timedelta_rhs or self.is_offset_rhs):
+            # datetime and timedelta/DateOffset
             if name not in ('__add__', '__radd__', '__sub__'):
                 raise TypeError("can only operate on a datetime with a rhs of "
                                 "a timedelta/DateOffset for addition and "
                                 "subtraction, but the operator [{name}] was "
                                 "passed".format(name=name))
 
-        elif (self.is_datetime_rhs and
-              (self.is_timedelta_lhs or self.is_offset_lhs)):
-            if name not in ('__add__', '__radd__', '__rsub__'):
-                raise TypeError("can only operate on a timedelta/DateOffset "
-                                "with a rhs of a datetime for addition, "
-                                "but the operator [{name}] was passed"
-                                .format(name=name))
-
-        # 2 datetimes
-        elif self.is_datetime_lhs and self.is_datetime_rhs:
-
+        elif self.is_datetime_rhs:
+            # 2 datetimes
             if name not in ('__sub__', '__rsub__'):
                 raise TypeError("can only operate on a datetimes for"
                                 " subtraction, but the operator [{name}] was"
@@ -445,18 +410,82 @@ def _validate(self, lvalues, rvalues, name):
                 raise ValueError("Incompatible tz's on datetime subtraction "
                                  "ops")
 
-        elif ((self.is_timedelta_lhs or self.is_offset_lhs) and
-              self.is_datetime_rhs):
+        else:
+            raise TypeError('cannot operate on a series without a rhs '
+                            'of a series/ndarray of type datetime64[ns] '
+                            'or a timedelta')
+
+    def _validate_timedelta(self, name):
+        # assumes self.is_timedelta_lhs
+
+        if self.is_integer_rhs or self.is_floating_rhs:
+            # timedelta and integer mul/div
+            self._check_timedelta_with_numeric(name)
+        elif self.is_timedelta_rhs or self.is_offset_rhs:
+            # 2 timedeltas
+            if name not in ('__div__', '__rdiv__', '__truediv__',
+                            '__rtruediv__', '__add__', '__radd__', '__sub__',
+                            '__rsub__'):
+                raise TypeError("can only operate on a timedeltas for addition"
+                                ", subtraction, and division, but the operator"
+                                " [{name}] was passed".format(name=name))
+        elif self.is_datetime_rhs:
+            if name not in ('__add__', '__radd__', '__rsub__'):
+                raise TypeError("can only operate on a timedelta/DateOffset "
+                                "with a rhs of a datetime for addition, "
+                                "but the operator [{name}] was passed"
+                                .format(name=name))
+        else:
+            raise TypeError('cannot operate on a series without a rhs '
+                            'of a series/ndarray of type datetime64[ns] '
+                            'or a timedelta')
+
+    def _validate_offset(self, name):
+        # assumes self.is_offset_lhs
+
+        if self.is_timedelta_rhs:
+            # 2 timedeltas
+            if name not in ('__div__', '__rdiv__', '__truediv__',
+                            '__rtruediv__', '__add__', '__radd__', '__sub__',
+                            '__rsub__'):
+                raise TypeError("can only operate on a timedeltas for addition"
+                                ", subtraction, and division, but the operator"
+                                " [{name}] was passed".format(name=name))
 
+        elif self.is_datetime_rhs:
             if name not in ('__add__', '__radd__'):
                 raise TypeError("can only operate on a timedelta/DateOffset "
                                 "and a datetime for addition, but the operator"
                                 " [{name}] was passed".format(name=name))
+
         else:
             raise TypeError('cannot operate on a series without a rhs '
                             'of a series/ndarray of type datetime64[ns] '
                             'or a timedelta')
 
+    def _validate(self, lvalues, rvalues, name):
+        if self.is_datetime_lhs:
+            return self._validate_datetime(lvalues, rvalues, name)
+        elif self.is_timedelta_lhs:
+            return self._validate_timedelta(name)
+        elif self.is_offset_lhs:
+            return self._validate_offset(name)
+
+        if ((self.is_integer_lhs or self.is_floating_lhs) and
+                self.is_timedelta_rhs):
+            self._check_timedelta_with_numeric(name)
+        else:
+            raise TypeError('cannot operate on a series without a rhs '
+                            'of a series/ndarray of type datetime64[ns] '
+                            'or a timedelta')
+
+    def _check_timedelta_with_numeric(self, name):
+        if name not in ('__div__', '__truediv__', '__mul__', '__rmul__'):
+            raise TypeError("can only operate on a timedelta and an "
+                            "integer or a float for division and "
+                            "multiplication, but the operator [{name}] "
+                            "was passed".format(name=name))
+
     def _convert_to_array(self, values, name=None, other=None):
         """converts values to ndarray"""
         from pandas.core.tools.timedeltas import to_timedelta

pandas/tests/series/test_operators.py

@@ -960,8 +960,51 @@ def test_timedelta64_ops_nat(self):
         assert_series_equal(timedelta_series / nan,
                             nat_series_dtype_timedelta)
 
+    @pytest.mark.parametrize('scalar_td', [timedelta(minutes=5, seconds=4),
+                                           Timedelta(minutes=5, seconds=4),
+                                           Timedelta('5m4s').to_timedelta64()])
+    def test_operators_timedelta64_with_timedelta(self, scalar_td):
+        # smoke tests
+        td1 = Series([timedelta(minutes=5, seconds=3)] * 3)
+        td1.iloc[2] = np.nan
+
+        td1 + scalar_td
+        scalar_td + td1
+        td1 - scalar_td
+        scalar_td - td1
+        td1 / scalar_td
+        scalar_td / td1
+
+    @pytest.mark.parametrize('scalar_td', [
+        timedelta(minutes=5, seconds=4),
+        pytest.param(Timedelta('5m4s'),
+                     marks=pytest.mark.xfail(reason="Timedelta.__floordiv__ "
+                                                    "bug GH#18846")),
+        Timedelta('5m4s').to_timedelta64()])
+    def test_operators_timedelta64_with_timedelta_invalid(self, scalar_td):
+        td1 = Series([timedelta(minutes=5, seconds=3)] * 3)
+        td1.iloc[2] = np.nan
+
+        # check that we are getting a TypeError
+        # with 'operate' (from core/ops.py) for the ops that are not
+        # defined
+        pattern = 'operate|unsupported|cannot'
+        with tm.assert_raises_regex(TypeError, pattern):
+            td1 * scalar_td
+        with tm.assert_raises_regex(TypeError, pattern):
+            scalar_td * td1
+        with tm.assert_raises_regex(TypeError, pattern):
+            td1 // scalar_td
+        with tm.assert_raises_regex(TypeError, pattern):
+            scalar_td // td1
+        with tm.assert_raises_regex(TypeError, pattern):
+            scalar_td ** td1
+        with tm.assert_raises_regex(TypeError, pattern):
+            td1 ** scalar_td
+
 
 class TestDatetimeSeriesArithmetic(object):
+
     def test_operators_datetimelike(self):
         def run_ops(ops, get_ser, test_ser):
 
@@ -976,16 +1019,6 @@ def run_ops(ops, get_ser, test_ser):
         # ## timedelta64 ###
         td1 = Series([timedelta(minutes=5, seconds=3)] * 3)
         td1.iloc[2] = np.nan
-        td2 = timedelta(minutes=5, seconds=4)
-        ops = ['__mul__', '__floordiv__', '__pow__', '__rmul__',
-               '__rfloordiv__', '__rpow__']
-        run_ops(ops, td1, td2)
-        td1 + td2
-        td2 + td1
-        td1 - td2
-        td2 - td1
-        td1 / td2
-        td2 / td1
 
         # ## datetime64 ###
         dt1 = Series([Timestamp('20111230'), Timestamp('20120101'),