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