Skip to content

Commit 893ac03

Browse files
committed
improve test for timestamp addition/subtraction
1 parent c407b73 commit 893ac03

File tree

1 file changed

+14
-6
lines changed

1 file changed

+14
-6
lines changed

pandas/tests/scalar/timestamp/test_arithmetic.py

Lines changed: 14 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -96,20 +96,28 @@ def test_addition_subtraction_types(self):
9696
assert type(ts + td64) == Timestamp
9797
assert type(ts - td64) == Timestamp
9898

99-
def test_addition_subtraction_preserve_frequency(self):
100-
ts = Timestamp('2014-03-05', freq='D')
101-
td = timedelta(days=1)
99+
@pytest.mark.parametrize('freq, td, td64', [
100+
('D', timedelta(days=1), np.timedelta64(1, 'D')),
101+
('W', timedelta(weeks=1), np.timedelta64(1, 'W')),
102+
('M', timedelta(months=1), np.timedelta64(1, 'M'))
103+
])
104+
def test_addition_subtraction_preserve_frequency(self, freq, td, td64):
105+
ts = Timestamp('2014-03-05', freq=freq)
102106
original_freq = ts.freq
103107

104108
with tm.assert_produces_warning(FutureWarning):
105109
# GH#22535 add/sub with integers is deprecated
106110
assert (ts + 1).freq == original_freq
107111
assert (ts - 1).freq == original_freq
108112

109-
assert (ts + td).freq == original_freq
110-
assert (ts - td).freq == original_freq
113+
assert (ts + 1 * original_freq).freq == original_freq
114+
assert (ts - 1 * original_freq).freq == original_freq
115+
116+
if td is not None:
117+
# timedelta does not support months as unit
118+
assert (ts + td).freq == original_freq
119+
assert (ts - td).freq == original_freq
111120

112-
td64 = np.timedelta64(1, 'D')
113121
assert (ts + td64).freq == original_freq
114122
assert (ts - td64).freq == original_freq
115123

0 commit comments

Comments
 (0)