ENH: TDA.__mul__ support non-nano (#47668)

* ENH: TDA.__mul__ support non-nano * fix freq

Author: jbrockmendel

pandas/core/arrays/timedeltas.py

@@ -398,7 +398,7 @@ def __mul__(self, other) -> TimedeltaArray:
  freq = None
  if self.freq is not None and not isna(other):
  freq = self.freq * other
- return type(self)(result, freq=freq)
+ return type(self)._simple_new(result, dtype=result.dtype, freq=freq)
 
  if not hasattr(other, "dtype"):
  # list, tuple
@@ -412,13 +412,14 @@ def __mul__(self, other) -> TimedeltaArray:
  # this multiplication will succeed only if all elements of other
  # are int or float scalars, so we will end up with
  # timedelta64[ns]-dtyped result
- result = [self[n] * other[n] for n in range(len(self))]
+ arr = self._ndarray
+ result = [arr[n] * other[n] for n in range(len(self))]
  result = np.array(result)
- return type(self)(result)
+ return type(self)._simple_new(result, dtype=result.dtype)
 
  # numpy will accept float or int dtype, raise TypeError for others
  result = self._ndarray * other
- return type(self)(result)
+ return type(self)._simple_new(result, dtype=result.dtype)
 
  __rmul__ = __mul__
 
@@ -446,7 +447,8 @@ def __truediv__(self, other):
  if self.freq is not None:
  # Tick division is not implemented, so operate on Timedelta
  freq = self.freq.delta / other
- return type(self)(result, freq=freq)
+ freq = to_offset(freq)
+ return type(self)._simple_new(result, dtype=result.dtype, freq=freq)
 
  if not hasattr(other, "dtype"):
  # e.g. list, tuple
@@ -462,6 +464,7 @@ def __truediv__(self, other):
  elif is_object_dtype(other.dtype):
  # We operate on raveled arrays to avoid problems in inference
  # on NaT
+ # TODO: tests with non-nano
  srav = self.ravel()
  orav = other.ravel()
  result_list = [srav[n] / orav[n] for n in range(len(srav))]
@@ -488,7 +491,7 @@ def __truediv__(self, other):
 
  else:
  result = self._ndarray / other
- return type(self)(result)
+ return type(self)._simple_new(result, dtype=result.dtype)
 
  @unpack_zerodim_and_defer("__rtruediv__")
  def __rtruediv__(self, other):

pandas/tests/arrays/test_timedeltas.py

@@ -1,3 +1,5 @@
+from datetime import timedelta
+
 import numpy as np
 import pytest
 
@@ -90,6 +92,53 @@ def test_add_pdnat(self, tda):
  assert result._reso == tda._reso
  assert result.isna().all()
 
+ def test_mul_scalar(self, tda):
+ other = 2
+ result = tda * other
+ expected = TimedeltaArray._simple_new(tda._ndarray * other, dtype=tda.dtype)
+ tm.assert_extension_array_equal(result, expected)
+ assert result._reso == tda._reso
+
+ def test_mul_listlike(self, tda):
+ other = np.arange(len(tda))
+ result = tda * other
+ expected = TimedeltaArray._simple_new(tda._ndarray * other, dtype=tda.dtype)
+ tm.assert_extension_array_equal(result, expected)
+ assert result._reso == tda._reso
+
+ def test_mul_listlike_object(self, tda):
+ other = np.arange(len(tda))
+ result = tda * other.astype(object)
+ expected = TimedeltaArray._simple_new(tda._ndarray * other, dtype=tda.dtype)
+ tm.assert_extension_array_equal(result, expected)
+ assert result._reso == tda._reso
+
+ def test_div_numeric_scalar(self, tda):
+ other = 2
+ result = tda / other
+ expected = TimedeltaArray._simple_new(tda._ndarray / other, dtype=tda.dtype)
+ tm.assert_extension_array_equal(result, expected)
+ assert result._reso == tda._reso
+
+ def test_div_td_scalar(self, tda):
+ other = timedelta(seconds=1)
+ result = tda / other
+ expected = tda._ndarray / np.timedelta64(1, "s")
+ tm.assert_numpy_array_equal(result, expected)
+
+ def test_div_numeric_array(self, tda):
+ other = np.arange(len(tda))
+ result = tda / other
+ expected = TimedeltaArray._simple_new(tda._ndarray / other, dtype=tda.dtype)
+ tm.assert_extension_array_equal(result, expected)
+ assert result._reso == tda._reso
+
+ def test_div_td_array(self, tda):
+ other = tda._ndarray + tda._ndarray[-1]
+ result = tda / other
+ expected = tda._ndarray / other
+ tm.assert_numpy_array_equal(result, expected)
+
 
 class TestTimedeltaArray:
  @pytest.mark.parametrize("dtype", [int, np.int32, np.int64, "uint32", "uint64"])