ENH: DatetimeIndex.snap support non-nano (#47537)

Author: jbrockmendel

pandas/core/indexes/datetimes.py

@@ -46,7 +46,6 @@
 from pandas.util._exceptions import find_stack_level
 
 from pandas.core.dtypes.common import (
-    DT64NS_DTYPE,
     is_datetime64_dtype,
     is_datetime64tz_dtype,
     is_scalar,
@@ -566,7 +565,7 @@ def snap(self, freq="S") -> DatetimeIndex:
         # Superdumb, punting on any optimizing
         freq = to_offset(freq)
 
-        snapped = np.empty(len(self), dtype=DT64NS_DTYPE)
+        dta = self._data.copy()
 
         for i, v in enumerate(self):
             s = v
@@ -577,9 +576,8 @@ def snap(self, freq="S") -> DatetimeIndex:
                     s = t0
                 else:
                     s = t1
-            snapped[i] = s
+            dta[i] = s
 
-        dta = DatetimeArray(snapped, dtype=self.dtype)
         return DatetimeIndex._simple_new(dta, name=self.name)
 
     # --------------------------------------------------------------------

pandas/tests/indexes/datetimes/methods/test_snap.py

@@ -7,10 +7,30 @@
 import pandas._testing as tm
 
 
+def astype_non_nano(dti_nano, unit):
+    # TODO(2.0): remove once DTI/DTA.astype supports non-nano
+    if unit == "ns":
+        return dti_nano
+
+    dta_nano = dti_nano._data
+    arr_nano = dta_nano._ndarray
+
+    arr = arr_nano.astype(f"M8[{unit}]")
+    if dti_nano.tz is None:
+        dtype = arr.dtype
+    else:
+        dtype = type(dti_nano.dtype)(tz=dti_nano.tz, unit=unit)
+    dta = type(dta_nano)._simple_new(arr, dtype=dtype)
+    dti = DatetimeIndex(dta, name=dti_nano.name)
+    assert dti.dtype == dtype
+    return dti
+
+
 @pytest.mark.filterwarnings("ignore::DeprecationWarning")
 @pytest.mark.parametrize("tz", [None, "Asia/Shanghai", "Europe/Berlin"])
 @pytest.mark.parametrize("name", [None, "my_dti"])
-def test_dti_snap(name, tz):
+@pytest.mark.parametrize("unit", ["ns", "us", "ms", "s"])
+def test_dti_snap(name, tz, unit):
     dti = DatetimeIndex(
         [
             "1/1/2002",
@@ -25,10 +45,12 @@ def test_dti_snap(name, tz):
         tz=tz,
         freq="D",
     )
+    dti = astype_non_nano(dti, unit)
 
     result = dti.snap(freq="W-MON")
     expected = date_range("12/31/2001", "1/7/2002", name=name, tz=tz, freq="w-mon")
     expected = expected.repeat([3, 4])
+    expected = astype_non_nano(expected, unit)
     tm.assert_index_equal(result, expected)
     assert result.tz == expected.tz
     assert result.freq is None
@@ -38,6 +60,7 @@ def test_dti_snap(name, tz):
 
     expected = date_range("1/1/2002", "1/7/2002", name=name, tz=tz, freq="b")
     expected = expected.repeat([1, 1, 1, 2, 2])
+    expected = astype_non_nano(expected, unit)
     tm.assert_index_equal(result, expected)
     assert result.tz == expected.tz
     assert result.freq is None