diff --git a/pandas/core/arrays/datetimelike.py b/pandas/core/arrays/datetimelike.py index 540442b7eaed4..df17388856117 100644 --- a/pandas/core/arrays/datetimelike.py +++ b/pandas/core/arrays/datetimelike.py @@ -36,7 +36,7 @@ ) from pandas.core.dtypes.generic import ABCDataFrame, ABCIndexClass, ABCSeries from pandas.core.dtypes.inference import is_array_like -from pandas.core.dtypes.missing import isna +from pandas.core.dtypes.missing import is_valid_nat_for_dtype, isna from pandas._typing import DatetimeLikeScalar from pandas.core import missing, nanops @@ -492,7 +492,10 @@ def __setitem__( elif isinstance(value, self._scalar_type): self._check_compatible_with(value) value = self._unbox_scalar(value) - elif isna(value) or value == iNaT: + elif is_valid_nat_for_dtype(value, self.dtype): + value = iNaT + elif not isna(value) and lib.is_integer(value) and value == iNaT: + # exclude misc e.g. object() and any NAs not allowed above value = iNaT else: msg = ( diff --git a/pandas/core/dtypes/missing.py b/pandas/core/dtypes/missing.py index f540e9297738a..6a681954fd902 100644 --- a/pandas/core/dtypes/missing.py +++ b/pandas/core/dtypes/missing.py @@ -559,3 +559,27 @@ def remove_na_arraylike(arr): return arr[notna(arr)] else: return arr[notna(lib.values_from_object(arr))] + + +def is_valid_nat_for_dtype(obj, dtype): + """ + isna check that excludes incompatible dtypes + + Parameters + ---------- + obj : object + dtype : np.datetime64, np.timedelta64, DatetimeTZDtype, or PeriodDtype + + Returns + ------- + bool + """ + if not isna(obj): + return False + if dtype.kind == "M": + return not isinstance(obj, np.timedelta64) + if dtype.kind == "m": + return not isinstance(obj, np.datetime64) + + # must be PeriodDType + return not isinstance(obj, (np.datetime64, np.timedelta64)) diff --git a/pandas/tests/arrays/test_datetimelike.py b/pandas/tests/arrays/test_datetimelike.py index 34fae1f4b1ab4..d9646feaf661e 100644 --- a/pandas/tests/arrays/test_datetimelike.py +++ b/pandas/tests/arrays/test_datetimelike.py @@ -651,3 +651,51 @@ def test_array_interface(self, period_index): result = np.asarray(arr, dtype="S20") expected = np.asarray(arr).astype("S20") tm.assert_numpy_array_equal(result, expected) + + +@pytest.mark.parametrize( + "array,casting_nats", + [ + ( + pd.TimedeltaIndex(["1 Day", "3 Hours", "NaT"])._data, + (pd.NaT, np.timedelta64("NaT", "ns")), + ), + ( + pd.date_range("2000-01-01", periods=3, freq="D")._data, + (pd.NaT, np.datetime64("NaT", "ns")), + ), + (pd.period_range("2000-01-01", periods=3, freq="D")._data, (pd.NaT,)), + ], + ids=lambda x: type(x).__name__, +) +def test_casting_nat_setitem_array(array, casting_nats): + expected = type(array)._from_sequence([pd.NaT, array[1], array[2]]) + + for nat in casting_nats: + arr = array.copy() + arr[0] = nat + tm.assert_equal(arr, expected) + + +@pytest.mark.parametrize( + "array,non_casting_nats", + [ + ( + pd.TimedeltaIndex(["1 Day", "3 Hours", "NaT"])._data, + (np.datetime64("NaT", "ns"),), + ), + ( + pd.date_range("2000-01-01", periods=3, freq="D")._data, + (np.timedelta64("NaT", "ns"),), + ), + ( + pd.period_range("2000-01-01", periods=3, freq="D")._data, + (np.datetime64("NaT", "ns"), np.timedelta64("NaT", "ns")), + ), + ], + ids=lambda x: type(x).__name__, +) +def test_invalid_nat_setitem_array(array, non_casting_nats): + for nat in non_casting_nats: + with pytest.raises(TypeError): + array[0] = nat