Skip to content

TST: parametrize and de-duplicate timedelta64 arithmetic tests #32091

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 8 commits into from
Feb 23, 2020
212 changes: 71 additions & 141 deletions pandas/tests/arithmetic/test_timedelta64.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,19 @@
get_upcast_box,
)


def assert_dtype(obj, expected_dtype):
"""
Helper to check the dtype for a Series, Index, or single-column DataFrame.
"""
if isinstance(obj, DataFrame):
dtype = obj.dtypes.iat[0]
else:
dtype = obj.dtype

assert dtype == expected_dtype


# ------------------------------------------------------------------
# Timedelta64[ns] dtype Comparisons

Expand Down Expand Up @@ -522,19 +535,35 @@ def test_tda_add_sub_index(self):
# -------------------------------------------------------------
# Binary operations TimedeltaIndex and timedelta-like

def test_tdi_iadd_timedeltalike(self, two_hours):
def test_tdi_iadd_timedeltalike(self, two_hours, box_with_array):
# only test adding/sub offsets as + is now numeric
rng = timedelta_range("1 days", "10 days")
expected = timedelta_range("1 days 02:00:00", "10 days 02:00:00", freq="D")

rng = tm.box_expected(rng, box_with_array)
expected = tm.box_expected(expected, box_with_array)

orig_rng = rng
rng += two_hours
tm.assert_index_equal(rng, expected)
tm.assert_equal(rng, expected)
if box_with_array is not pd.Index:
# Check that operation is actually inplace
tm.assert_equal(orig_rng, expected)

def test_tdi_isub_timedeltalike(self, two_hours):
def test_tdi_isub_timedeltalike(self, two_hours, box_with_array):
# only test adding/sub offsets as - is now numeric
rng = timedelta_range("1 days", "10 days")
expected = timedelta_range("0 days 22:00:00", "9 days 22:00:00")

rng = tm.box_expected(rng, box_with_array)
expected = tm.box_expected(expected, box_with_array)

orig_rng = rng
rng -= two_hours
tm.assert_index_equal(rng, expected)
tm.assert_equal(rng, expected)
if box_with_array is not pd.Index:
# Check that operation is actually inplace
tm.assert_equal(orig_rng, expected)

# -------------------------------------------------------------

Expand Down Expand Up @@ -1013,15 +1042,6 @@ def test_td64arr_add_datetime64_nat(self, box_with_array):
# ------------------------------------------------------------------
# Invalid __add__/__sub__ operations

# TODO: moved from frame tests; needs parametrization/de-duplication
def test_td64_df_add_int_frame(self):
# GH#22696 Check that we don't dispatch to numpy implementation,
# which treats int64 as m8[ns]
tdi = pd.timedelta_range("1", periods=3)
df = tdi.to_frame()
other = pd.DataFrame([1, 2, 3], index=tdi) # indexed like `df`
assert_invalid_addsub_type(df, other)

@pytest.mark.parametrize("pi_freq", ["D", "W", "Q", "H"])
@pytest.mark.parametrize("tdi_freq", [None, "H"])
def test_td64arr_sub_periodlike(self, box_with_array, tdi_freq, pi_freq):
Expand Down Expand Up @@ -1100,6 +1120,9 @@ def test_td64arr_add_sub_int(self, box_with_array, one):

def test_td64arr_add_sub_integer_array(self, box_with_array):
# GH#19959, deprecated GH#22535
# GH#22696 for DataFrame case, check that we don't dispatch to numpy
# implementation, which treats int64 as m8[ns]

rng = timedelta_range("1 days 09:00:00", freq="H", periods=3)
tdarr = tm.box_expected(rng, box_with_array)
other = tm.box_expected([4, 3, 2], box_with_array)
Expand All @@ -1119,60 +1142,6 @@ def test_td64arr_addsub_integer_array_no_freq(self, box_with_array):
# ------------------------------------------------------------------
# Operations with timedelta-like others

# TODO: this was taken from tests.series.test_ops; de-duplicate
def test_operators_timedelta64_with_timedelta(self, scalar_td):
# smoke tests
td1 = Series([timedelta(minutes=5, seconds=3)] * 3)
td1.iloc[2] = np.nan

td1 + scalar_td
scalar_td + td1
td1 - scalar_td
scalar_td - td1
td1 / scalar_td
scalar_td / td1

# TODO: this was taken from tests.series.test_ops; de-duplicate
def test_timedelta64_operations_with_timedeltas(self):
# td operate with td
td1 = Series([timedelta(minutes=5, seconds=3)] * 3)
td2 = timedelta(minutes=5, seconds=4)
result = td1 - td2
expected = Series([timedelta(seconds=0)] * 3) - Series(
[timedelta(seconds=1)] * 3
)
assert result.dtype == "m8[ns]"
tm.assert_series_equal(result, expected)

result2 = td2 - td1
expected = Series([timedelta(seconds=1)] * 3) - Series(
[timedelta(seconds=0)] * 3
)
tm.assert_series_equal(result2, expected)

# roundtrip
tm.assert_series_equal(result + td2, td1)

# Now again, using pd.to_timedelta, which should build
# a Series or a scalar, depending on input.
td1 = Series(pd.to_timedelta(["00:05:03"] * 3))
td2 = pd.to_timedelta("00:05:04")
result = td1 - td2
expected = Series([timedelta(seconds=0)] * 3) - Series(
[timedelta(seconds=1)] * 3
)
assert result.dtype == "m8[ns]"
tm.assert_series_equal(result, expected)

result2 = td2 - td1
expected = Series([timedelta(seconds=1)] * 3) - Series(
[timedelta(seconds=0)] * 3
)
tm.assert_series_equal(result2, expected)

# roundtrip
tm.assert_series_equal(result + td2, td1)

def test_td64arr_add_td64_array(self, box_with_array):
box = box_with_array
dti = pd.date_range("2016-01-01", periods=3)
Expand Down Expand Up @@ -1203,7 +1172,6 @@ def test_td64arr_sub_td64_array(self, box_with_array):
result = tdarr - tdi
tm.assert_equal(result, expected)

# TODO: parametrize over [add, sub, radd, rsub]?
@pytest.mark.parametrize(
"names",
[
Expand Down Expand Up @@ -1232,17 +1200,11 @@ def test_td64arr_add_sub_tdi(self, box, names):

result = tdi + ser
tm.assert_equal(result, expected)
if box is not pd.DataFrame:
assert result.dtype == "timedelta64[ns]"
else:
assert result.dtypes[0] == "timedelta64[ns]"
assert_dtype(result, "timedelta64[ns]")

result = ser + tdi
tm.assert_equal(result, expected)
if box is not pd.DataFrame:
assert result.dtype == "timedelta64[ns]"
else:
assert result.dtypes[0] == "timedelta64[ns]"
assert_dtype(result, "timedelta64[ns]")

expected = Series(
[Timedelta(hours=-3), Timedelta(days=1, hours=-4)], name=names[2]
Expand All @@ -1251,17 +1213,11 @@ def test_td64arr_add_sub_tdi(self, box, names):

result = tdi - ser
tm.assert_equal(result, expected)
if box is not pd.DataFrame:
assert result.dtype == "timedelta64[ns]"
else:
assert result.dtypes[0] == "timedelta64[ns]"
assert_dtype(result, "timedelta64[ns]")

result = ser - tdi
tm.assert_equal(result, -expected)
if box is not pd.DataFrame:
assert result.dtype == "timedelta64[ns]"
else:
assert result.dtypes[0] == "timedelta64[ns]"
assert_dtype(result, "timedelta64[ns]")

def test_td64arr_add_sub_td64_nat(self, box_with_array):
# GH#23320 special handling for timedelta64("NaT")
Expand Down Expand Up @@ -1296,6 +1252,7 @@ def test_td64arr_sub_NaT(self, box_with_array):

def test_td64arr_add_timedeltalike(self, two_hours, box_with_array):
# only test adding/sub offsets as + is now numeric
# GH#10699 for Tick cases
box = box_with_array
rng = timedelta_range("1 days", "10 days")
expected = timedelta_range("1 days 02:00:00", "10 days 02:00:00", freq="D")
Expand All @@ -1305,8 +1262,12 @@ def test_td64arr_add_timedeltalike(self, two_hours, box_with_array):
result = rng + two_hours
tm.assert_equal(result, expected)

result = two_hours + rng
tm.assert_equal(result, expected)

def test_td64arr_sub_timedeltalike(self, two_hours, box_with_array):
# only test adding/sub offsets as - is now numeric
# GH#10699 for Tick cases
box = box_with_array
rng = timedelta_range("1 days", "10 days")
expected = timedelta_range("0 days 22:00:00", "9 days 22:00:00")
Expand All @@ -1317,46 +1278,12 @@ def test_td64arr_sub_timedeltalike(self, two_hours, box_with_array):
result = rng - two_hours
tm.assert_equal(result, expected)

result = two_hours - rng
tm.assert_equal(result, -expected)

# ------------------------------------------------------------------
# __add__/__sub__ with DateOffsets and arrays of DateOffsets

# TODO: this was taken from tests.series.test_operators; de-duplicate
def test_timedelta64_operations_with_DateOffset(self):
# GH#10699
td = Series([timedelta(minutes=5, seconds=3)] * 3)
result = td + pd.offsets.Minute(1)
expected = Series([timedelta(minutes=6, seconds=3)] * 3)
tm.assert_series_equal(result, expected)

result = td - pd.offsets.Minute(1)
expected = Series([timedelta(minutes=4, seconds=3)] * 3)
tm.assert_series_equal(result, expected)

with tm.assert_produces_warning(PerformanceWarning):
result = td + Series(
[pd.offsets.Minute(1), pd.offsets.Second(3), pd.offsets.Hour(2)]
)
expected = Series(
[
timedelta(minutes=6, seconds=3),
timedelta(minutes=5, seconds=6),
timedelta(hours=2, minutes=5, seconds=3),
]
)
tm.assert_series_equal(result, expected)

result = td + pd.offsets.Minute(1) + pd.offsets.Second(12)
expected = Series([timedelta(minutes=6, seconds=15)] * 3)
tm.assert_series_equal(result, expected)

# valid DateOffsets
for do in ["Hour", "Minute", "Second", "Day", "Micro", "Milli", "Nano"]:
op = getattr(pd.offsets, do)
td + op(5)
op(5) + td
td - op(5)
op(5) - td

@pytest.mark.parametrize(
"names", [(None, None, None), ("foo", "bar", None), ("foo", "foo", "foo")]
)
Expand Down Expand Up @@ -1561,26 +1488,6 @@ class TestTimedeltaArraylikeMulDivOps:
# Tests for timedelta64[ns]
# __mul__, __rmul__, __div__, __rdiv__, __floordiv__, __rfloordiv__

# TODO: Moved from tests.series.test_operators; needs cleanup
@pytest.mark.parametrize("m", [1, 3, 10])
@pytest.mark.parametrize("unit", ["D", "h", "m", "s", "ms", "us", "ns"])
def test_timedelta64_conversions(self, m, unit):
startdate = Series(pd.date_range("2013-01-01", "2013-01-03"))
enddate = Series(pd.date_range("2013-03-01", "2013-03-03"))

ser = enddate - startdate
ser[2] = np.nan

# op
expected = Series([x / np.timedelta64(m, unit) for x in ser])
result = ser / np.timedelta64(m, unit)
tm.assert_series_equal(result, expected)

# reverse op
expected = Series([Timedelta(np.timedelta64(m, unit)) / x for x in ser])
result = np.timedelta64(m, unit) / ser
tm.assert_series_equal(result, expected)

# ------------------------------------------------------------------
# Multiplication
# organized with scalar others first, then array-like
Expand Down Expand Up @@ -1734,6 +1641,29 @@ def test_td64arr_div_tdlike_scalar(self, two_hours, box_with_array):
expected = 1 / expected
tm.assert_equal(result, expected)

@pytest.mark.parametrize("m", [1, 3, 10])
@pytest.mark.parametrize("unit", ["D", "h", "m", "s", "ms", "us", "ns"])
def test_td64arr_div_td64_scalar(self, m, unit, box_with_array):
startdate = Series(pd.date_range("2013-01-01", "2013-01-03"))
enddate = Series(pd.date_range("2013-03-01", "2013-03-03"))

ser = enddate - startdate
ser[2] = np.nan
flat = ser
ser = tm.box_expected(ser, box_with_array)

# op
expected = Series([x / np.timedelta64(m, unit) for x in flat])
expected = tm.box_expected(expected, box_with_array)
result = ser / np.timedelta64(m, unit)
tm.assert_equal(result, expected)

# reverse op
expected = Series([Timedelta(np.timedelta64(m, unit)) / x for x in flat])
expected = tm.box_expected(expected, box_with_array)
result = np.timedelta64(m, unit) / ser
tm.assert_equal(result, expected)

def test_td64arr_div_tdlike_scalar_with_nat(self, two_hours, box_with_array):
rng = TimedeltaIndex(["1 days", pd.NaT, "2 days"], name="foo")
expected = pd.Float64Index([12, np.nan, 24], name="foo")
Expand Down