REF: stronger typing in _box_func (#46917)

Author: jbrockmendel

pandas/core/arrays/datetimes.py

@@ -538,13 +538,10 @@ def _check_compatible_with(self, other, setitem: bool = False):
     # -----------------------------------------------------------------
     # Descriptive Properties
 
-    def _box_func(self, x) -> Timestamp | NaTType:
-        if isinstance(x, np.datetime64):
-            # GH#42228
-            # Argument 1 to "signedinteger" has incompatible type "datetime64";
-            # expected "Union[SupportsInt, Union[str, bytes], SupportsIndex]"
-            x = np.int64(x)  # type: ignore[arg-type]
-        ts = Timestamp(x, tz=self.tz)
+    def _box_func(self, x: np.datetime64) -> Timestamp | NaTType:
+        # GH#42228
+        value = x.view("i8")
+        ts = Timestamp(value, tz=self.tz)
         # Non-overlapping identity check (left operand type: "Timestamp",
         # right operand type: "NaTType")
         if ts is not NaT:  # type: ignore[comparison-overlap]

pandas/core/arrays/timedeltas.py

@@ -154,7 +154,7 @@ class TimedeltaArray(dtl.TimelikeOps):
     # Note: ndim must be defined to ensure NaT.__richcmp__(TimedeltaArray)
     #  operates pointwise.
 
-    def _box_func(self, x) -> Timedelta | NaTType:
+    def _box_func(self, x: np.timedelta64) -> Timedelta | NaTType:
         return Timedelta(x, unit="ns")
 
     @property

pandas/core/nanops.py

@@ -16,7 +16,6 @@
 from pandas._libs import (
     NaT,
     NaTType,
-    Timedelta,
     iNaT,
     lib,
 )
@@ -367,19 +366,23 @@ def _wrap_results(result, dtype: np.dtype, fill_value=None):
                 result = np.datetime64("NaT", "ns")
             else:
                 result = np.int64(result).view("datetime64[ns]")
+            # retain original unit
+            result = result.astype(dtype, copy=False)
         else:
             # If we have float dtype, taking a view will give the wrong result
             result = result.astype(dtype)
     elif is_timedelta64_dtype(dtype):
         if not isinstance(result, np.ndarray):
-            if result == fill_value:
-                result = np.nan
+            if result == fill_value or np.isnan(result):
+                result = np.timedelta64("NaT").astype(dtype)
 
-            # raise if we have a timedelta64[ns] which is too large
-            if np.fabs(result) > lib.i8max:
+            elif np.fabs(result) > lib.i8max:
+                # raise if we have a timedelta64[ns] which is too large
                 raise ValueError("overflow in timedelta operation")
+            else:
+                # return a timedelta64 with the original unit
+                result = np.int64(result).astype(dtype, copy=False)
 
-            result = Timedelta(result, unit="ns")
         else:
             result = result.astype("m8[ns]").view(dtype)
 
@@ -641,7 +644,7 @@ def _mask_datetimelike_result(
         result[axis_mask] = iNaT  # type: ignore[index]
     else:
         if mask.any():
-            return NaT
+            return np.int64(iNaT).view(orig_values.dtype)
     return result
 
 

pandas/tests/arrays/timedeltas/test_reductions.py

@@ -147,7 +147,7 @@ def test_std(self, add):
 
         if getattr(arr, "tz", None) is None:
             result = nanops.nanstd(np.asarray(arr), skipna=True)
-            assert isinstance(result, Timedelta)
+            assert isinstance(result, np.timedelta64)
             assert result == expected
 
         result = arr.std(skipna=False)
@@ -158,7 +158,8 @@ def test_std(self, add):
 
         if getattr(arr, "tz", None) is None:
             result = nanops.nanstd(np.asarray(arr), skipna=False)
-            assert result is pd.NaT
+            assert isinstance(result, np.timedelta64)
+            assert np.isnat(result)
 
     def test_median(self):
         tdi = pd.TimedeltaIndex(["0H", "3H", "NaT", "5H06m", "0H", "2H"])

pandas/tests/test_nanops.py

@@ -1020,7 +1020,8 @@ def test_nanmean_skipna_false(self, dtype):
         arr[-1, -1] = "NaT"
 
         result = nanops.nanmean(arr, skipna=False)
-        assert result is pd.NaT
+        assert np.isnat(result)
+        assert result.dtype == dtype
 
         result = nanops.nanmean(arr, axis=0, skipna=False)
         expected = np.array([4, 5, "NaT"], dtype=arr.dtype)