Skip to content

Commit

Permalink
TST: parametrize over dt64 unit (#55818)
Browse files Browse the repository at this point in the history
  • Loading branch information
jbrockmendel authored Nov 4, 2023
1 parent c2cdeaf commit 407ec33
Show file tree
Hide file tree
Showing 18 changed files with 179 additions and 120 deletions.
12 changes: 12 additions & 0 deletions pandas/_testing/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
"""
Expand Down Expand Up @@ -1121,6 +1132,7 @@ def shares_memory(left, right) -> bool:
"getitem",
"get_locales",
"getMixedTypeDict",
"get_finest_unit",
"get_obj",
"get_op_from_name",
"getPeriodData",
Expand Down
3 changes: 3 additions & 0 deletions pandas/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -1303,6 +1303,9 @@ def unit(request):
return request.param


unit2 = unit


# ----------------------------------------------------------------
# Dtypes
# ----------------------------------------------------------------
Expand Down
30 changes: 16 additions & 14 deletions pandas/tests/arithmetic/test_datetime64.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 = [
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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)
Expand Down
8 changes: 2 additions & 6 deletions pandas/tests/arrays/test_datetimes.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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
Expand Down
9 changes: 1 addition & 8 deletions pandas/tests/frame/methods/test_reset_index.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"])
Expand Down
32 changes: 22 additions & 10 deletions pandas/tests/indexes/datetimes/methods/test_delete.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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(
Expand All @@ -84,7 +96,7 @@ def test_delete_slice(self):
],
freq=None,
name="idx",
)
).as_unit(unit)

cases = {
(0, 1, 2): expected_0_2,
Expand Down
33 changes: 21 additions & 12 deletions pandas/tests/indexes/datetimes/methods/test_insert.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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),
Expand All @@ -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),
Expand All @@ -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),
Expand Down
7 changes: 6 additions & 1 deletion pandas/tests/indexes/datetimes/methods/test_repeat.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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(
[
Expand Down
5 changes: 5 additions & 0 deletions pandas/tests/indexes/datetimes/methods/test_round.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand All @@ -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"])
Expand Down
Loading

0 comments on commit 407ec33

Please sign in to comment.