From a4654dd1d12f605c7e9e09bc548a0451988b4710 Mon Sep 17 00:00:00 2001
From: Brock <jbrockmendel@gmail.com>
Date: Tue, 28 Jun 2022 14:42:50 -0700
Subject: [PATCH] ENH: DatetimeIndex.snap support non-nano

---
 pandas/core/indexes/datetimes.py              |  6 ++---
 .../indexes/datetimes/methods/test_snap.py    | 25 ++++++++++++++++++-
 2 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/pandas/core/indexes/datetimes.py b/pandas/core/indexes/datetimes.py
index 4bc01b8f4ddb0..aacc2a2225743 100644
--- a/pandas/core/indexes/datetimes.py
+++ b/pandas/core/indexes/datetimes.py
@@ -44,7 +44,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,
@@ -552,7 +551,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
@@ -563,9 +562,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)
 
     # --------------------------------------------------------------------
diff --git a/pandas/tests/indexes/datetimes/methods/test_snap.py b/pandas/tests/indexes/datetimes/methods/test_snap.py
index e591441c4f148..a94d00d919082 100644
--- a/pandas/tests/indexes/datetimes/methods/test_snap.py
+++ b/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