TST: parametrize over unit (#55806)

Author: jbrockmendel

pandas/conftest.py

@@ -1295,6 +1295,14 @@ def utc_fixture(request):
 utc_fixture2 = utc_fixture
 
 
+@pytest.fixture(params=["s", "ms", "us", "ns"])
+def unit(request):
+    """
+    datetime64 units we support.
+    """
+    return request.param
+
+
 # ----------------------------------------------------------------
 # Dtypes
 # ----------------------------------------------------------------

pandas/tests/arithmetic/test_datetime64.py

@@ -987,21 +987,21 @@ def test_dt64arr_sub_timestamp_tzaware(self, box_with_array):
         tm.assert_equal(ser - ts, expected)
         tm.assert_equal(ts - ser, -expected)
 
-    def test_dt64arr_sub_NaT(self, box_with_array):
+    def test_dt64arr_sub_NaT(self, box_with_array, unit):
         # GH#18808
-        dti = DatetimeIndex([NaT, Timestamp("19900315")])
+        dti = DatetimeIndex([NaT, Timestamp("19900315")]).as_unit(unit)
         ser = tm.box_expected(dti, box_with_array)
 
         result = ser - NaT
-        expected = Series([NaT, NaT], dtype="timedelta64[ns]")
+        expected = Series([NaT, NaT], dtype=f"timedelta64[{unit}]")
         expected = tm.box_expected(expected, box_with_array)
         tm.assert_equal(result, expected)
 
         dti_tz = dti.tz_localize("Asia/Tokyo")
         ser_tz = tm.box_expected(dti_tz, box_with_array)
 
         result = ser_tz - NaT
-        expected = Series([NaT, NaT], dtype="timedelta64[ns]")
+        expected = Series([NaT, NaT], dtype=f"timedelta64[{unit}]")
         expected = tm.box_expected(expected, box_with_array)
         tm.assert_equal(result, expected)
 
@@ -1872,15 +1872,15 @@ def test_operators_datetimelike_invalid(
             # Smoke test
             op(arg2)
 
-    def test_sub_single_tz(self):
+    def test_sub_single_tz(self, unit):
         # GH#12290
-        s1 = Series([Timestamp("2016-02-10", tz="America/Sao_Paulo")])
-        s2 = Series([Timestamp("2016-02-08", tz="America/Sao_Paulo")])
+        s1 = Series([Timestamp("2016-02-10", tz="America/Sao_Paulo")]).dt.as_unit(unit)
+        s2 = Series([Timestamp("2016-02-08", tz="America/Sao_Paulo")]).dt.as_unit(unit)
         result = s1 - s2
-        expected = Series([Timedelta("2days")])
+        expected = Series([Timedelta("2days")]).dt.as_unit(unit)
         tm.assert_series_equal(result, expected)
         result = s2 - s1
-        expected = Series([Timedelta("-2days")])
+        expected = Series([Timedelta("-2days")]).dt.as_unit(unit)
         tm.assert_series_equal(result, expected)
 
     def test_dt64tz_series_sub_dtitz(self):
@@ -1895,13 +1895,19 @@ def test_dt64tz_series_sub_dtitz(self):
         res = ser - dti
         tm.assert_series_equal(res, expected)
 
-    def test_sub_datetime_compat(self):
+    def test_sub_datetime_compat(self, unit):
         # see GH#14088
-        s = Series([datetime(2016, 8, 23, 12, tzinfo=pytz.utc), NaT])
+        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 = Series([Timedelta("1 days"), NaT])
-        tm.assert_series_equal(s - dt, exp)
-        tm.assert_series_equal(s - Timestamp(dt), exp)
+        exp_unit = unit
+        if unit in ["s", "ms"]:
+            # The datetime object has "us" so we upcast
+            exp_unit = "us"
+        exp = Series([Timedelta("1 days"), NaT]).dt.as_unit(exp_unit)
+        result = ser - dt
+        tm.assert_series_equal(result, exp)
+        result2 = ser - Timestamp(dt)
+        tm.assert_series_equal(result2, exp)
 
     def test_dt64_series_add_mixed_tick_DateOffset(self):
         # GH#4532
@@ -1922,11 +1928,11 @@ def test_dt64_series_add_mixed_tick_DateOffset(self):
         )
         tm.assert_series_equal(result, expected)
 
-    def test_datetime64_ops_nat(self):
+    def test_datetime64_ops_nat(self, unit):
         # GH#11349
-        datetime_series = Series([NaT, Timestamp("19900315")])
-        nat_series_dtype_timestamp = Series([NaT, NaT], dtype="datetime64[ns]")
-        single_nat_dtype_datetime = Series([NaT], dtype="datetime64[ns]")
+        datetime_series = Series([NaT, Timestamp("19900315")]).dt.as_unit(unit)
+        nat_series_dtype_timestamp = Series([NaT, NaT], dtype=f"datetime64[{unit}]")
+        single_nat_dtype_datetime = Series([NaT], dtype=f"datetime64[{unit}]")
 
         # subtraction
         tm.assert_series_equal(-NaT + datetime_series, nat_series_dtype_timestamp)
@@ -2097,16 +2103,16 @@ def test_dti_sub_tdi(self, tz_naive_fixture):
         with pytest.raises(TypeError, match=msg):
             tdi.values - dti
 
-    def test_dti_isub_tdi(self, tz_naive_fixture):
+    def test_dti_isub_tdi(self, tz_naive_fixture, unit):
         # GH#17558
         tz = tz_naive_fixture
-        dti = DatetimeIndex([Timestamp("2017-01-01", tz=tz)] * 10)
-        tdi = pd.timedelta_range("0 days", periods=10)
-        expected = date_range("2017-01-01", periods=10, tz=tz, freq="-1D")
+        dti = DatetimeIndex([Timestamp("2017-01-01", tz=tz)] * 10).as_unit(unit)
+        tdi = pd.timedelta_range("0 days", periods=10, unit=unit)
+        expected = date_range("2017-01-01", periods=10, tz=tz, freq="-1D", unit=unit)
         expected = expected._with_freq(None)
 
         # isub with TimedeltaIndex
-        result = DatetimeIndex([Timestamp("2017-01-01", tz=tz)] * 10)
+        result = DatetimeIndex([Timestamp("2017-01-01", tz=tz)] * 10).as_unit(unit)
         result -= tdi
         tm.assert_index_equal(result, expected)
 
@@ -2124,7 +2130,7 @@ def test_dti_isub_tdi(self, tz_naive_fixture):
             tdi -= dti
 
         # isub with timedelta64 array
-        result = DatetimeIndex([Timestamp("2017-01-01", tz=tz)] * 10)
+        result = DatetimeIndex([Timestamp("2017-01-01", tz=tz)] * 10).as_unit(unit)
         result -= tdi.values
         tm.assert_index_equal(result, expected)
 
@@ -2158,13 +2164,13 @@ def test_dta_add_sub_index(self, tz_naive_fixture):
         expected = dti - tdi
         tm.assert_index_equal(result, expected)
 
-    def test_sub_dti_dti(self):
+    def test_sub_dti_dti(self, unit):
         # previously performed setop (deprecated in 0.16.0), now changed to
         # return subtraction -> TimeDeltaIndex (GH ...)
 
-        dti = date_range("20130101", periods=3)
-        dti_tz = date_range("20130101", periods=3).tz_localize("US/Eastern")
-        expected = TimedeltaIndex([0, 0, 0])
+        dti = date_range("20130101", periods=3, unit=unit)
+        dti_tz = date_range("20130101", periods=3, unit=unit).tz_localize("US/Eastern")
+        expected = TimedeltaIndex([0, 0, 0]).as_unit(unit)
 
         result = dti - dti
         tm.assert_index_equal(result, expected)
@@ -2183,16 +2189,16 @@ def test_sub_dti_dti(self):
         tm.assert_index_equal(dti, expected)
 
         # different length raises ValueError
-        dti1 = date_range("20130101", periods=3)
-        dti2 = date_range("20130101", periods=4)
+        dti1 = date_range("20130101", periods=3, unit=unit)
+        dti2 = date_range("20130101", periods=4, unit=unit)
         msg = "cannot add indices of unequal length"
         with pytest.raises(ValueError, match=msg):
             dti1 - dti2
 
         # NaN propagation
-        dti1 = DatetimeIndex(["2012-01-01", np.nan, "2012-01-03"])
-        dti2 = DatetimeIndex(["2012-01-02", "2012-01-03", np.nan])
-        expected = TimedeltaIndex(["1 days", np.nan, np.nan])
+        dti1 = DatetimeIndex(["2012-01-01", np.nan, "2012-01-03"]).as_unit(unit)
+        dti2 = DatetimeIndex(["2012-01-02", "2012-01-03", np.nan]).as_unit(unit)
+        expected = TimedeltaIndex(["1 days", np.nan, np.nan]).as_unit(unit)
         result = dti2 - dti1
         tm.assert_index_equal(result, expected)
 
@@ -2295,17 +2301,17 @@ def test_ops_nat_mixed_datetime64_timedelta64(self):
             nat_series_dtype_timestamp,
         )
 
-    def test_ufunc_coercions(self):
-        idx = date_range("2011-01-01", periods=3, freq="2D", name="x")
+    def test_ufunc_coercions(self, unit):
+        idx = date_range("2011-01-01", periods=3, freq="2D", name="x", unit=unit)
 
         delta = np.timedelta64(1, "D")
-        exp = date_range("2011-01-02", periods=3, freq="2D", name="x")
+        exp = date_range("2011-01-02", periods=3, freq="2D", name="x", unit=unit)
         for result in [idx + delta, np.add(idx, delta)]:
             assert isinstance(result, DatetimeIndex)
             tm.assert_index_equal(result, exp)
             assert result.freq == "2D"
 
-        exp = date_range("2010-12-31", periods=3, freq="2D", name="x")
+        exp = date_range("2010-12-31", periods=3, freq="2D", name="x", unit=unit)
 
         for result in [idx - delta, np.subtract(idx, delta)]:
             assert isinstance(result, DatetimeIndex)
@@ -2318,13 +2324,17 @@ def test_ufunc_coercions(self):
         delta = np.array(
             [np.timedelta64(1, "D"), np.timedelta64(2, "D"), np.timedelta64(3, "D")]
         )
-        exp = DatetimeIndex(["2011-01-02", "2011-01-05", "2011-01-08"], name="x")
+        exp = DatetimeIndex(
+            ["2011-01-02", "2011-01-05", "2011-01-08"], name="x"
+        ).as_unit(unit)
 
         for result in [idx + delta, np.add(idx, delta)]:
             tm.assert_index_equal(result, exp)
             assert result.freq == exp.freq
 
-        exp = DatetimeIndex(["2010-12-31", "2011-01-01", "2011-01-02"], name="x")
+        exp = DatetimeIndex(
+            ["2010-12-31", "2011-01-01", "2011-01-02"], name="x"
+        ).as_unit(unit)
         for result in [idx - delta, np.subtract(idx, delta)]:
             assert isinstance(result, DatetimeIndex)
             tm.assert_index_equal(result, exp)

pandas/tests/arithmetic/test_timedelta64.py

@@ -336,17 +336,17 @@ def test_subtraction_ops(self):
 
         result = tdi - td
         expected = TimedeltaIndex(["0 days", NaT, "1 days"], name="foo")
-        tm.assert_index_equal(result, expected, check_names=False)
+        tm.assert_index_equal(result, expected)
 
         result = td - tdi
         expected = TimedeltaIndex(["0 days", NaT, "-1 days"], name="foo")
-        tm.assert_index_equal(result, expected, check_names=False)
+        tm.assert_index_equal(result, expected)
 
         result = dti - td
         expected = DatetimeIndex(
             ["20121231", "20130101", "20130102"], freq="D", name="bar"
         )
-        tm.assert_index_equal(result, expected, check_names=False)
+        tm.assert_index_equal(result, expected)
 
         result = dt - tdi
         expected = DatetimeIndex(["20121231", NaT, "20121230"], name="foo")

pandas/tests/base/test_conversion.py

@@ -141,35 +141,41 @@ def test_categorial_datetimelike(self, method):
         result = method(i)[0]
         assert isinstance(result, Timestamp)
 
-    def test_iter_box(self):
+    def test_iter_box_dt64(self, unit):
         vals = [Timestamp("2011-01-01"), Timestamp("2011-01-02")]
-        s = Series(vals)
-        assert s.dtype == "datetime64[ns]"
-        for res, exp in zip(s, vals):
+        ser = Series(vals).dt.as_unit(unit)
+        assert ser.dtype == f"datetime64[{unit}]"
+        for res, exp in zip(ser, vals):
             assert isinstance(res, Timestamp)
             assert res.tz is None
             assert res == exp
+            assert res.unit == unit
 
+    def test_iter_box_dt64tz(self, unit):
         vals = [
             Timestamp("2011-01-01", tz="US/Eastern"),
             Timestamp("2011-01-02", tz="US/Eastern"),
         ]
-        s = Series(vals)
+        ser = Series(vals).dt.as_unit(unit)
 
-        assert s.dtype == "datetime64[ns, US/Eastern]"
-        for res, exp in zip(s, vals):
+        assert ser.dtype == f"datetime64[{unit}, US/Eastern]"
+        for res, exp in zip(ser, vals):
             assert isinstance(res, Timestamp)
             assert res.tz == exp.tz
             assert res == exp
+            assert res.unit == unit
 
+    def test_iter_box_timedelta64(self, unit):
         # timedelta
         vals = [Timedelta("1 days"), Timedelta("2 days")]
-        s = Series(vals)
-        assert s.dtype == "timedelta64[ns]"
-        for res, exp in zip(s, vals):
+        ser = Series(vals).dt.as_unit(unit)
+        assert ser.dtype == f"timedelta64[{unit}]"
+        for res, exp in zip(ser, vals):
             assert isinstance(res, Timedelta)
             assert res == exp
+            assert res.unit == unit
 
+    def test_iter_box_period(self):
         # period
         vals = [pd.Period("2011-01-01", freq="M"), pd.Period("2011-01-02", freq="M")]
         s = Series(vals)
@@ -370,7 +376,7 @@ def test_to_numpy_copy(arr, as_series):
 
 
 @pytest.mark.parametrize("as_series", [True, False])
-def test_to_numpy_dtype(as_series):
+def test_to_numpy_dtype(as_series, unit):
     tz = "US/Eastern"
     obj = pd.DatetimeIndex(["2000", "2001"], tz=tz)
     if as_series: