REF: dispatch TDBlock._can_hold_element to TimedeltaArray._validate_setitem_value (#38674)

Author: jbrockmendel

pandas/core/generic.py

@@ -8858,7 +8858,7 @@ def _where(
                     elif len(cond[icond]) == len(other):
 
                         # try to not change dtype at first
-                        new_other = np.asarray(self)
+                        new_other = self._values
                         new_other = new_other.copy()
                         new_other[icond] = other
                         other = new_other

pandas/core/internals/blocks.py

@@ -1,4 +1,3 @@
-from datetime import timedelta
 import inspect
 import re
 from typing import TYPE_CHECKING, Any, List, Optional, Type, Union, cast
@@ -8,7 +7,6 @@
 
 from pandas._libs import (
     Interval,
-    NaT,
     Period,
     Timestamp,
     algos as libalgos,
@@ -86,6 +84,7 @@
 
 if TYPE_CHECKING:
     from pandas import Index
+    from pandas.core.arrays._mixins import NDArrayBackedExtensionArray
 
 
 class Block(PandasObject):
@@ -919,7 +918,11 @@ def setitem(self, indexer, value):
         if self._can_hold_element(value):
             # We only get here for non-Extension Blocks, so _try_coerce_args
             #  is only relevant for DatetimeBlock and TimedeltaBlock
-            if lib.is_scalar(value):
+            if self.dtype.kind in ["m", "M"]:
+                arr = self.array_values().T
+                arr[indexer] = value
+                return self
+            elif lib.is_scalar(value):
                 value = convert_scalar_for_putitemlike(value, values.dtype)
 
         else:
@@ -1064,6 +1067,17 @@ def putmask(
         if self._can_hold_element(new):
             # We only get here for non-Extension Blocks, so _try_coerce_args
             #  is only relevant for DatetimeBlock and TimedeltaBlock
+            if self.dtype.kind in ["m", "M"]:
+                blk = self
+                if not inplace:
+                    blk = self.copy()
+                arr = blk.array_values()
+                arr = cast("NDArrayBackedExtensionArray", arr)
+                if transpose:
+                    arr = arr.T
+                arr.putmask(mask, new)
+                return [blk]
+
             if lib.is_scalar(new):
                 new = convert_scalar_for_putitemlike(new, self.values.dtype)
 
@@ -2376,16 +2390,6 @@ def _maybe_coerce_values(self, values):
     def _holder(self):
         return TimedeltaArray
 
-    def _can_hold_element(self, element: Any) -> bool:
-        tipo = maybe_infer_dtype_type(element)
-        if tipo is not None:
-            return issubclass(tipo.type, np.timedelta64)
-        elif element is NaT:
-            return True
-        elif isinstance(element, (timedelta, np.timedelta64)):
-            return True
-        return is_valid_nat_for_dtype(element, self.dtype)
-
     def fillna(self, value, **kwargs):
         # TODO(EA2D): if we operated on array_values, TDA.fillna would handle
         #  raising here.

pandas/tests/indexing/test_indexing.py

@@ -10,7 +10,7 @@
 from pandas.core.dtypes.common import is_float_dtype, is_integer_dtype
 
 import pandas as pd
-from pandas import DataFrame, Index, NaT, Series, date_range
+from pandas import DataFrame, Index, NaT, Series, date_range, offsets, timedelta_range
 import pandas._testing as tm
 from pandas.core.indexing import maybe_numeric_slice, non_reducing_slice
 from pandas.tests.indexing.common import _mklbl
@@ -970,19 +970,22 @@ class TestDatetimelikeCoercion:
     @pytest.mark.parametrize("indexer", [setitem, loc, iloc])
     def test_setitem_dt64_string_scalar(self, tz_naive_fixture, indexer):
         # dispatching _can_hold_element to underling DatetimeArray
-        # TODO(EA2D) use tz_naive_fixture once DatetimeBlock is backed by DTA
         tz = tz_naive_fixture
 
         dti = date_range("2016-01-01", periods=3, tz=tz)
         ser = Series(dti)
 
         values = ser._values
 
-        indexer(ser)[0] = "2018-01-01"
+        newval = "2018-01-01"
+        values._validate_setitem_value(newval)
+
+        indexer(ser)[0] = newval
 
         if tz is None:
             # TODO(EA2D): we can make this no-copy in tz-naive case too
             assert ser.dtype == dti.dtype
+            assert ser._values._data is values._data
         else:
             assert ser._values is values
 
@@ -993,7 +996,6 @@ def test_setitem_dt64_string_scalar(self, tz_naive_fixture, indexer):
     @pytest.mark.parametrize("indexer", [setitem, loc, iloc])
     def test_setitem_dt64_string_values(self, tz_naive_fixture, indexer, key, box):
         # dispatching _can_hold_element to underling DatetimeArray
-        # TODO(EA2D) use tz_naive_fixture once DatetimeBlock is backed by DTA
         tz = tz_naive_fixture
 
         if isinstance(key, slice) and indexer is loc:
@@ -1012,9 +1014,44 @@ def test_setitem_dt64_string_values(self, tz_naive_fixture, indexer, key, box):
         if tz is None:
             # TODO(EA2D): we can make this no-copy in tz-naive case too
             assert ser.dtype == dti.dtype
+            assert ser._values._data is values._data
         else:
             assert ser._values is values
 
+    @pytest.mark.parametrize("scalar", ["3 Days", offsets.Hour(4)])
+    @pytest.mark.parametrize("indexer", [setitem, loc, iloc])
+    def test_setitem_td64_scalar(self, indexer, scalar):
+        # dispatching _can_hold_element to underling TimedeltaArray
+        tdi = timedelta_range("1 Day", periods=3)
+        ser = Series(tdi)
+
+        values = ser._values
+        values._validate_setitem_value(scalar)
+
+        indexer(ser)[0] = scalar
+        assert ser._values._data is values._data
+
+    @pytest.mark.parametrize("box", [list, np.array, pd.array])
+    @pytest.mark.parametrize(
+        "key", [[0, 1], slice(0, 2), np.array([True, True, False])]
+    )
+    @pytest.mark.parametrize("indexer", [setitem, loc, iloc])
+    def test_setitem_td64_string_values(self, indexer, key, box):
+        # dispatching _can_hold_element to underling TimedeltaArray
+        if isinstance(key, slice) and indexer is loc:
+            key = slice(0, 1)
+
+        tdi = timedelta_range("1 Day", periods=3)
+        ser = Series(tdi)
+
+        values = ser._values
+
+        newvals = box(["10 Days", "44 hours"])
+        values._validate_setitem_value(newvals)
+
+        indexer(ser)[key] = newvals
+        assert ser._values._data is values._data
+
 
 def test_extension_array_cross_section():
     # A cross-section of a homogeneous EA should be an EA