TST: parametrize over dt64 unit (#55818)

Author: jbrockmendel

pandas/_testing/__init__.py

@@ -1019,6 +1019,17 @@ def iat(x):
 
 # -----------------------------------------------------------------------------
 
+_UNITS = ["s", "ms", "us", "ns"]
+
+
+def get_finest_unit(left: str, right: str):
+    """
+    Find the higher of two datetime64 units.
+    """
+    if _UNITS.index(left) >= _UNITS.index(right):
+        return left
+    return right
+
 
 def shares_memory(left, right) -> bool:
     """
@@ -1121,6 +1132,7 @@ def shares_memory(left, right) -> bool:
     "getitem",
     "get_locales",
     "getMixedTypeDict",
+    "get_finest_unit",
     "get_obj",
     "get_op_from_name",
     "getPeriodData",

pandas/conftest.py

@@ -1303,6 +1303,9 @@ def unit(request):
     return request.param
 
 
+unit2 = unit
+
+
 # ----------------------------------------------------------------
 # Dtypes
 # ----------------------------------------------------------------

pandas/tests/arithmetic/test_datetime64.py

@@ -1696,8 +1696,8 @@ def test_dt64_series_arith_overflow(self):
         tm.assert_series_equal(res, -expected)
 
     def test_datetimeindex_sub_timestamp_overflow(self):
-        dtimax = pd.to_datetime(["2021-12-28 17:19", Timestamp.max])
-        dtimin = pd.to_datetime(["2021-12-28 17:19", Timestamp.min])
+        dtimax = pd.to_datetime(["2021-12-28 17:19", Timestamp.max]).as_unit("ns")
+        dtimin = pd.to_datetime(["2021-12-28 17:19", Timestamp.min]).as_unit("ns")
 
         tsneg = Timestamp("1950-01-01").as_unit("ns")
         ts_neg_variants = [
@@ -1735,11 +1735,11 @@ def test_datetimeindex_sub_timestamp_overflow(self):
 
     def test_datetimeindex_sub_datetimeindex_overflow(self):
         # GH#22492, GH#22508
-        dtimax = pd.to_datetime(["2021-12-28 17:19", Timestamp.max])
-        dtimin = pd.to_datetime(["2021-12-28 17:19", Timestamp.min])
+        dtimax = pd.to_datetime(["2021-12-28 17:19", Timestamp.max]).as_unit("ns")
+        dtimin = pd.to_datetime(["2021-12-28 17:19", Timestamp.min]).as_unit("ns")
 
-        ts_neg = pd.to_datetime(["1950-01-01", "1950-01-01"])
-        ts_pos = pd.to_datetime(["1980-01-01", "1980-01-01"])
+        ts_neg = pd.to_datetime(["1950-01-01", "1950-01-01"]).as_unit("ns")
+        ts_pos = pd.to_datetime(["1980-01-01", "1980-01-01"]).as_unit("ns")
 
         # General tests
         expected = Timestamp.max._value - ts_pos[1]._value
@@ -1815,12 +1815,16 @@ def test_operators_datetimelike(self):
         td1 + dt1
         dt1 + td1
 
-    def test_dt64ser_sub_datetime_dtype(self):
+    def test_dt64ser_sub_datetime_dtype(self, unit):
         ts = Timestamp(datetime(1993, 1, 7, 13, 30, 00))
         dt = datetime(1993, 6, 22, 13, 30)
-        ser = Series([ts])
-        result = pd.to_timedelta(np.abs(ser - dt))
-        assert result.dtype == "timedelta64[ns]"
+        ser = Series([ts], dtype=f"M8[{unit}]")
+        result = ser - dt
+
+        # the expected unit is the max of `unit` and the unit imputed to `dt`,
+        #  which is "us"
+        exp_unit = tm.get_finest_unit(unit, "us")
+        assert result.dtype == f"timedelta64[{exp_unit}]"
 
     # -------------------------------------------------------------
     # TODO: This next block of tests came from tests.series.test_operators,
@@ -1899,10 +1903,8 @@ def test_sub_datetime_compat(self, unit):
         # see GH#14088
         ser = Series([datetime(2016, 8, 23, 12, tzinfo=pytz.utc), NaT]).dt.as_unit(unit)
         dt = datetime(2016, 8, 22, 12, tzinfo=pytz.utc)
-        exp_unit = unit
-        if unit in ["s", "ms"]:
-            # The datetime object has "us" so we upcast
-            exp_unit = "us"
+        # The datetime object has "us" so we upcast lower units
+        exp_unit = tm.get_finest_unit(unit, "us")
         exp = Series([Timedelta("1 days"), NaT]).dt.as_unit(exp_unit)
         result = ser - dt
         tm.assert_series_equal(result, exp)

pandas/tests/arrays/test_datetimes.py

@@ -15,10 +15,7 @@
 import numpy as np
 import pytest
 
-from pandas._libs.tslibs import (
-    npy_unit_to_abbrev,
-    tz_compare,
-)
+from pandas._libs.tslibs import tz_compare
 
 from pandas.core.dtypes.dtypes import DatetimeTZDtype
 
@@ -235,8 +232,7 @@ def test_add_timedeltalike_scalar_mismatched_reso(self, dta_dti, scalar):
         dta, dti = dta_dti
 
         td = pd.Timedelta(scalar)
-        exp_reso = max(dta._creso, td._creso)
-        exp_unit = npy_unit_to_abbrev(exp_reso)
+        exp_unit = tm.get_finest_unit(dta.unit, td.unit)
 
         expected = (dti + td)._data.as_unit(exp_unit)
         result = dta + scalar

pandas/tests/frame/methods/test_reset_index.py

@@ -76,19 +76,12 @@ def test_reset_index_tz(self, tz_aware_fixture):
 
         expected = DataFrame(
             {
-                "idx": [
-                    datetime(2011, 1, 1),
-                    datetime(2011, 1, 2),
-                    datetime(2011, 1, 3),
-                    datetime(2011, 1, 4),
-                    datetime(2011, 1, 5),
-                ],
+                "idx": idx,
                 "a": range(5),
                 "b": ["A", "B", "C", "D", "E"],
             },
             columns=["idx", "a", "b"],
         )
-        expected["idx"] = expected["idx"].apply(lambda d: Timestamp(d, tz=tz))
         tm.assert_frame_equal(df.reset_index(), expected)
 
     @pytest.mark.parametrize("tz", ["US/Eastern", "dateutil/US/Eastern"])

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

@@ -9,19 +9,25 @@
 
 
 class TestDelete:
-    def test_delete(self):
-        idx = date_range(start="2000-01-01", periods=5, freq="ME", name="idx")
+    def test_delete(self, unit):
+        idx = date_range(
+            start="2000-01-01", periods=5, freq="ME", name="idx", unit=unit
+        )
 
         # preserve freq
-        expected_0 = date_range(start="2000-02-01", periods=4, freq="ME", name="idx")
-        expected_4 = date_range(start="2000-01-01", periods=4, freq="ME", name="idx")
+        expected_0 = date_range(
+            start="2000-02-01", periods=4, freq="ME", name="idx", unit=unit
+        )
+        expected_4 = date_range(
+            start="2000-01-01", periods=4, freq="ME", name="idx", unit=unit
+        )
 
         # reset freq to None
         expected_1 = DatetimeIndex(
             ["2000-01-31", "2000-03-31", "2000-04-30", "2000-05-31"],
             freq=None,
             name="idx",
-        )
+        ).as_unit(unit)
 
         cases = {
             0: expected_0,
@@ -64,12 +70,18 @@ def test_delete2(self, tz):
         assert result.freqstr == "h"
         assert result.tz == expected.tz
 
-    def test_delete_slice(self):
-        idx = date_range(start="2000-01-01", periods=10, freq="D", name="idx")
+    def test_delete_slice(self, unit):
+        idx = date_range(
+            start="2000-01-01", periods=10, freq="D", name="idx", unit=unit
+        )
 
         # preserve freq
-        expected_0_2 = date_range(start="2000-01-04", periods=7, freq="D", name="idx")
-        expected_7_9 = date_range(start="2000-01-01", periods=7, freq="D", name="idx")
+        expected_0_2 = date_range(
+            start="2000-01-04", periods=7, freq="D", name="idx", unit=unit
+        )
+        expected_7_9 = date_range(
+            start="2000-01-01", periods=7, freq="D", name="idx", unit=unit
+        )
 
         # reset freq to None
         expected_3_5 = DatetimeIndex(
@@ -84,7 +96,7 @@ def test_delete_slice(self):
             ],
             freq=None,
             name="idx",
-        )
+        ).as_unit(unit)
 
         cases = {
             (0, 1, 2): expected_0_2,

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

@@ -52,13 +52,15 @@ def test_insert_empty_preserves_freq(self, tz_naive_fixture):
         result = dti.insert(0, item)
         assert result.freq is None
 
-    def test_insert(self):
-        idx = DatetimeIndex(["2000-01-04", "2000-01-01", "2000-01-02"], name="idx")
+    def test_insert(self, unit):
+        idx = DatetimeIndex(
+            ["2000-01-04", "2000-01-01", "2000-01-02"], name="idx"
+        ).as_unit(unit)
 
         result = idx.insert(2, datetime(2000, 1, 5))
         exp = DatetimeIndex(
             ["2000-01-04", "2000-01-01", "2000-01-05", "2000-01-02"], name="idx"
-        )
+        ).as_unit(unit)
         tm.assert_index_equal(result, exp)
 
         # insertion of non-datetime should coerce to object index
@@ -76,31 +78,32 @@ def test_insert(self):
         tm.assert_index_equal(result, expected)
         assert result.name == expected.name
 
-        idx = date_range("1/1/2000", periods=3, freq="ME", name="idx")
+    def test_insert2(self, unit):
+        idx = date_range("1/1/2000", periods=3, freq="ME", name="idx", unit=unit)
 
         # preserve freq
         expected_0 = DatetimeIndex(
             ["1999-12-31", "2000-01-31", "2000-02-29", "2000-03-31"],
             name="idx",
             freq="ME",
-        )
+        ).as_unit(unit)
         expected_3 = DatetimeIndex(
             ["2000-01-31", "2000-02-29", "2000-03-31", "2000-04-30"],
             name="idx",
             freq="ME",
-        )
+        ).as_unit(unit)
 
         # reset freq to None
         expected_1_nofreq = DatetimeIndex(
             ["2000-01-31", "2000-01-31", "2000-02-29", "2000-03-31"],
             name="idx",
             freq=None,
-        )
+        ).as_unit(unit)
         expected_3_nofreq = DatetimeIndex(
             ["2000-01-31", "2000-02-29", "2000-03-31", "2000-01-02"],
             name="idx",
             freq=None,
-        )
+        ).as_unit(unit)
 
         cases = [
             (0, datetime(1999, 12, 31), expected_0),
@@ -116,22 +119,28 @@ def test_insert(self):
             assert result.name == expected.name
             assert result.freq == expected.freq
 
+    def test_insert3(self, unit):
+        idx = date_range("1/1/2000", periods=3, freq="ME", name="idx", unit=unit)
+
         # reset freq to None
         result = idx.insert(3, datetime(2000, 1, 2))
         expected = DatetimeIndex(
             ["2000-01-31", "2000-02-29", "2000-03-31", "2000-01-02"],
             name="idx",
             freq=None,
-        )
+        ).as_unit(unit)
         tm.assert_index_equal(result, expected)
         assert result.name == expected.name
         assert result.freq is None
 
+    def test_insert4(self, unit):
         for tz in ["US/Pacific", "Asia/Singapore"]:
-            idx = date_range("1/1/2000 09:00", periods=6, freq="h", tz=tz, name="idx")
+            idx = date_range(
+                "1/1/2000 09:00", periods=6, freq="h", tz=tz, name="idx", unit=unit
+            )
             # preserve freq
             expected = date_range(
-                "1/1/2000 09:00", periods=7, freq="h", tz=tz, name="idx"
+                "1/1/2000 09:00", periods=7, freq="h", tz=tz, name="idx", unit=unit
             )
             for d in [
                 Timestamp("2000-01-01 15:00", tz=tz),
@@ -156,7 +165,7 @@ def test_insert(self):
                 name="idx",
                 tz=tz,
                 freq=None,
-            )
+            ).as_unit(unit)
             # reset freq to None
             for d in [
                 Timestamp("2000-01-01 10:00", tz=tz),

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

@@ -11,13 +11,14 @@
 
 class TestRepeat:
     def test_repeat_range(self, tz_naive_fixture):
-        tz = tz_naive_fixture
         rng = date_range("1/1/2000", "1/1/2001")
 
         result = rng.repeat(5)
         assert result.freq is None
         assert len(result) == 5 * len(rng)
 
+    def test_repeat_range2(self, tz_naive_fixture):
+        tz = tz_naive_fixture
         index = date_range("2001-01-01", periods=2, freq="D", tz=tz)
         exp = DatetimeIndex(
             ["2001-01-01", "2001-01-01", "2001-01-02", "2001-01-02"], tz=tz
@@ -26,6 +27,8 @@ def test_repeat_range(self, tz_naive_fixture):
             tm.assert_index_equal(res, exp)
             assert res.freq is None
 
+    def test_repeat_range3(self, tz_naive_fixture):
+        tz = tz_naive_fixture
         index = date_range("2001-01-01", periods=2, freq="2D", tz=tz)
         exp = DatetimeIndex(
             ["2001-01-01", "2001-01-01", "2001-01-03", "2001-01-03"], tz=tz
@@ -34,6 +37,8 @@ def test_repeat_range(self, tz_naive_fixture):
             tm.assert_index_equal(res, exp)
             assert res.freq is None
 
+    def test_repeat_range4(self, tz_naive_fixture):
+        tz = tz_naive_fixture
         index = DatetimeIndex(["2001-01-01", "NaT", "2003-01-01"], tz=tz)
         exp = DatetimeIndex(
             [

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

@@ -71,6 +71,8 @@ def test_round(self, tz_naive_fixture):
         with pytest.raises(ValueError, match=msg):
             elt.round(freq="ME")
 
+    def test_round2(self, tz_naive_fixture):
+        tz = tz_naive_fixture
         # GH#14440 & GH#15578
         index = DatetimeIndex(["2016-10-17 12:00:00.0015"], tz=tz)
         result = index.round("ms")
@@ -80,11 +82,14 @@ def test_round(self, tz_naive_fixture):
         for freq in ["us", "ns"]:
             tm.assert_index_equal(index, index.round(freq))
 
+    def test_round3(self, tz_naive_fixture):
+        tz = tz_naive_fixture
         index = DatetimeIndex(["2016-10-17 12:00:00.00149"], tz=tz)
         result = index.round("ms")
         expected = DatetimeIndex(["2016-10-17 12:00:00.001000"], tz=tz)
         tm.assert_index_equal(result, expected)
 
+    def test_round4(self, tz_naive_fixture):
         index = DatetimeIndex(["2016-10-17 12:00:00.001501031"])
         result = index.round("10ns")
         expected = DatetimeIndex(["2016-10-17 12:00:00.001501030"])