Skip to content

Commit

Permalink
TST: parametrize over unit (#55806)
Browse files Browse the repository at this point in the history
  • Loading branch information
jbrockmendel authored Nov 3, 2023
1 parent 0cdb37c commit 8e0411f
Show file tree
Hide file tree
Showing 9 changed files with 162 additions and 135 deletions.
8 changes: 8 additions & 0 deletions pandas/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
# ----------------------------------------------------------------
Expand Down
88 changes: 49 additions & 39 deletions pandas/tests/arithmetic/test_datetime64.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down Expand Up @@ -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):
Expand All @@ -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
Expand All @@ -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)
Expand Down Expand Up @@ -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)

Expand All @@ -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)

Expand Down Expand Up @@ -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)
Expand All @@ -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)

Expand Down Expand Up @@ -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)
Expand All @@ -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)
Expand Down
6 changes: 3 additions & 3 deletions pandas/tests/arithmetic/test_timedelta64.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down
28 changes: 17 additions & 11 deletions pandas/tests/base/test_conversion.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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:
Expand Down
Loading

0 comments on commit 8e0411f

Please sign in to comment.