ENH: GH4521 A Series of dtype timedelta64[ns] can now be divided/multipled by an integer series

Author: jreback

doc/source/release.rst

@@ -53,6 +53,8 @@ pandas 0.13
   - Add ``rename`` and ``set_names`` methods to ``Index`` as well as
     ``set_names``, ``set_levels``, ``set_labels`` to ``MultiIndex``.
     (:issue:`4039`)
+  - A Series of dtype ``Timedelta64[ns]`` can now be divided/multiplied
+    by an integer series (:issue`4521`)
 
 **API Changes**
 

pandas/core/series.py

@@ -91,6 +91,7 @@ def wrapper(self, other, name=name):
 
         is_timedelta_lhs = com.is_timedelta64_dtype(self)
         is_datetime_lhs  = com.is_datetime64_dtype(self)
+        is_integer_lhs   = lvalues.dtype == np.int64
 
         if is_datetime_lhs or is_timedelta_lhs:
 
@@ -115,8 +116,8 @@ def convert_to_array(values):
                     # py3 compat where dtype is 'm' but is an integer
                     if values.dtype.kind == 'm':
                         values = values.astype('timedelta64[ns]')
-                    else:
-                        raise ValueError("incompatible type for a datetime/timedelta operation")
+                    elif name not in ['__div__','__mul__']:
+                        raise TypeError("incompatible type for a datetime/timedelta operation")
                 elif isinstance(values[0],DateOffset):
                     # handle DateOffsets
                     values = pa.array([ v.delta for v in values ])
@@ -131,9 +132,22 @@ def convert_to_array(values):
 
             is_datetime_rhs  = com.is_datetime64_dtype(rvalues)
             is_timedelta_rhs = com.is_timedelta64_dtype(rvalues) or (not is_datetime_rhs and _np_version_under1p7)
+            is_integer_rhs = rvalues.dtype == np.int64
+            mask = None
+
+            # timedelta and integer mul/div
+            if (is_timedelta_lhs and is_integer_rhs) or (is_integer_lhs and is_timedelta_rhs):
+
+                if name not in ['__div__','__mul__']:
+                    raise TypeError("can only operate on a timedelta and an integer for "
+                                    "division, but the operator [%s] was passed" % name)
+                dtype = 'timedelta64[ns]'
+                mask = isnull(lvalues) | isnull(rvalues)
+                lvalues = lvalues.astype(np.int64)
+                rvalues = rvalues.astype(np.int64)
 
             # 2 datetimes or 2 timedeltas
-            if (is_timedelta_lhs and is_timedelta_rhs) or (is_datetime_lhs and
+            elif (is_timedelta_lhs and is_timedelta_rhs) or (is_datetime_lhs and
                     is_datetime_rhs):
                 if is_datetime_lhs and name != '__sub__':
                     raise TypeError("can only operate on a datetimes for subtraction, "
@@ -143,13 +157,7 @@ def convert_to_array(values):
                                     "addition and subtraction, but the operator [%s] was passed" % name)
 
                 dtype = 'timedelta64[ns]'
-
                 mask = isnull(lvalues) | isnull(rvalues)
-                if mask.any():
-                    def wrap_results(x):
-                        x = pa.array(x,dtype='timedelta64[ns]')
-                        np.putmask(x,mask,tslib.iNaT)
-                        return x
 
             # datetime and timedelta
             elif is_timedelta_rhs and is_datetime_lhs:
@@ -167,9 +175,18 @@ def wrap_results(x):
                 dtype = 'M8[ns]'
 
             else:
-                raise ValueError('cannot operate on a series with out a rhs '
-                                 'of a series/ndarray of type datetime64[ns] '
-                                 'or a timedelta')
+                raise TypeError('cannot operate on a series with out a rhs '
+                                'of a series/ndarray of type datetime64[ns] '
+                                'or a timedelta')
+
+            # if we need to mask the results
+            if mask is not None:
+                if mask.any():
+                    def f(x):
+                        x = pa.array(x,dtype='timedelta64[ns]')
+                        np.putmask(x,mask,tslib.iNaT)
+                        return x
+                    wrap_results = f
 
             lvalues = lvalues.view('i8')
             rvalues = rvalues.view('i8')

pandas/tests/test_series.py

@@ -1999,6 +1999,51 @@ def test_operators_timedelta64(self):
                 [Timestamp('20130101 9:01:00.005'), Timestamp('20130101 9:02:00.005')])
             assert_series_equal(result, expected)
 
+        # GH 4521
+        # divide/multiply by integers
+        startdate = Series(date_range('2013-01-01', '2013-01-03'))
+        enddate = Series(date_range('2013-03-01', '2013-03-03'))
+
+        s1 = enddate - startdate
+        s1[2] = np.nan
+        s2 = Series([2, 3, 4])
+        expected = Series(s1.values.astype(np.int64) / s2, dtype='m8[ns]')
+        expected[2] = np.nan
+        result = s1 / s2
+        assert_series_equal(result,expected)
+
+        s2 = Series([20, 30, 40])
+        expected = Series(s1.values.astype(np.int64) / s2, dtype='m8[ns]')
+        expected[2] = np.nan
+        result = s1 / s2
+        assert_series_equal(result,expected)
+
+        result = s1 / 2
+        expected = Series(s1.values.astype(np.int64) / 2, dtype='m8[ns]')
+        expected[2] = np.nan
+        assert_series_equal(result,expected)
+
+        s2 = Series([20, 30, 40])
+        expected = Series(s1.values.astype(np.int64) * s2, dtype='m8[ns]')
+        expected[2] = np.nan
+        result = s1 * s2
+        assert_series_equal(result,expected)
+
+        result = s1 * 2
+        expected = Series(s1.values.astype(np.int64) * 2, dtype='m8[ns]')
+        expected[2] = np.nan
+        assert_series_equal(result,expected)
+
+        self.assertRaises(TypeError, s1.__div__, s2.astype(float))
+        self.assertRaises(TypeError, s1.__mul__, s2.astype(float))
+        self.assertRaises(TypeError, s1.__div__, 2.)
+        self.assertRaises(TypeError, s1.__mul__, 2.)
+
+        self.assertRaises(TypeError, s1.__add__, 1)
+        self.assertRaises(TypeError, s1.__sub__, 1)
+        self.assertRaises(TypeError, s1.__add__, s2.values)
+        self.assertRaises(TypeError, s1.__sub__, s2.values)
+
     def test_timedelta64_equal_timedelta_supported_ops(self):
         ser = Series([Timestamp('20130301'), Timestamp('20130228 23:00:00'),
                       Timestamp('20130228 22:00:00'),