test_arithmetic.py 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265
  1. # -*- coding: utf-8 -*-
  2. from datetime import timedelta
  3. import numpy as np
  4. import pytest
  5. from pandas.errors import NullFrequencyError
  6. import pandas as pd
  7. from pandas import Timedelta, TimedeltaIndex, timedelta_range
  8. import pandas.util.testing as tm
  9. @pytest.fixture(params=[pd.offsets.Hour(2), timedelta(hours=2),
  10. np.timedelta64(2, 'h'), Timedelta(hours=2)],
  11. ids=str)
  12. def delta(request):
  13. # Several ways of representing two hours
  14. return request.param
  15. @pytest.fixture(params=['B', 'D'])
  16. def freq(request):
  17. return request.param
  18. class TestTimedeltaIndexArithmetic(object):
  19. # Addition and Subtraction Operations
  20. # -------------------------------------------------------------
  21. # TimedeltaIndex.shift is used by __add__/__sub__
  22. def test_tdi_shift_empty(self):
  23. # GH#9903
  24. idx = pd.TimedeltaIndex([], name='xxx')
  25. tm.assert_index_equal(idx.shift(0, freq='H'), idx)
  26. tm.assert_index_equal(idx.shift(3, freq='H'), idx)
  27. def test_tdi_shift_hours(self):
  28. # GH#9903
  29. idx = pd.TimedeltaIndex(['5 hours', '6 hours', '9 hours'], name='xxx')
  30. tm.assert_index_equal(idx.shift(0, freq='H'), idx)
  31. exp = pd.TimedeltaIndex(['8 hours', '9 hours', '12 hours'], name='xxx')
  32. tm.assert_index_equal(idx.shift(3, freq='H'), exp)
  33. exp = pd.TimedeltaIndex(['2 hours', '3 hours', '6 hours'], name='xxx')
  34. tm.assert_index_equal(idx.shift(-3, freq='H'), exp)
  35. def test_tdi_shift_minutes(self):
  36. # GH#9903
  37. idx = pd.TimedeltaIndex(['5 hours', '6 hours', '9 hours'], name='xxx')
  38. tm.assert_index_equal(idx.shift(0, freq='T'), idx)
  39. exp = pd.TimedeltaIndex(['05:03:00', '06:03:00', '9:03:00'],
  40. name='xxx')
  41. tm.assert_index_equal(idx.shift(3, freq='T'), exp)
  42. exp = pd.TimedeltaIndex(['04:57:00', '05:57:00', '8:57:00'],
  43. name='xxx')
  44. tm.assert_index_equal(idx.shift(-3, freq='T'), exp)
  45. def test_tdi_shift_int(self):
  46. # GH#8083
  47. trange = pd.to_timedelta(range(5), unit='d') + pd.offsets.Hour(1)
  48. result = trange.shift(1)
  49. expected = TimedeltaIndex(['1 days 01:00:00', '2 days 01:00:00',
  50. '3 days 01:00:00',
  51. '4 days 01:00:00', '5 days 01:00:00'],
  52. freq='D')
  53. tm.assert_index_equal(result, expected)
  54. def test_tdi_shift_nonstandard_freq(self):
  55. # GH#8083
  56. trange = pd.to_timedelta(range(5), unit='d') + pd.offsets.Hour(1)
  57. result = trange.shift(3, freq='2D 1s')
  58. expected = TimedeltaIndex(['6 days 01:00:03', '7 days 01:00:03',
  59. '8 days 01:00:03', '9 days 01:00:03',
  60. '10 days 01:00:03'], freq='D')
  61. tm.assert_index_equal(result, expected)
  62. def test_shift_no_freq(self):
  63. # GH#19147
  64. tdi = TimedeltaIndex(['1 days 01:00:00', '2 days 01:00:00'], freq=None)
  65. with pytest.raises(NullFrequencyError):
  66. tdi.shift(2)
  67. # -------------------------------------------------------------
  68. # Binary operations TimedeltaIndex and integer
  69. def test_tdi_add_int(self, one):
  70. # Variants of `one` for #19012
  71. rng = timedelta_range('1 days 09:00:00', freq='H', periods=10)
  72. with tm.assert_produces_warning(FutureWarning, check_stacklevel=False):
  73. # GH#22535
  74. result = rng + one
  75. expected = timedelta_range('1 days 10:00:00', freq='H', periods=10)
  76. tm.assert_index_equal(result, expected)
  77. def test_tdi_iadd_int(self, one):
  78. rng = timedelta_range('1 days 09:00:00', freq='H', periods=10)
  79. expected = timedelta_range('1 days 10:00:00', freq='H', periods=10)
  80. with tm.assert_produces_warning(FutureWarning, check_stacklevel=False):
  81. # GH#22535
  82. rng += one
  83. tm.assert_index_equal(rng, expected)
  84. def test_tdi_sub_int(self, one):
  85. rng = timedelta_range('1 days 09:00:00', freq='H', periods=10)
  86. with tm.assert_produces_warning(FutureWarning, check_stacklevel=False):
  87. # GH#22535
  88. result = rng - one
  89. expected = timedelta_range('1 days 08:00:00', freq='H', periods=10)
  90. tm.assert_index_equal(result, expected)
  91. def test_tdi_isub_int(self, one):
  92. rng = timedelta_range('1 days 09:00:00', freq='H', periods=10)
  93. expected = timedelta_range('1 days 08:00:00', freq='H', periods=10)
  94. with tm.assert_produces_warning(FutureWarning, check_stacklevel=False):
  95. # GH#22535
  96. rng -= one
  97. tm.assert_index_equal(rng, expected)
  98. # -------------------------------------------------------------
  99. # __add__/__sub__ with integer arrays
  100. @pytest.mark.parametrize('box', [np.array, pd.Index])
  101. def test_tdi_add_integer_array(self, box):
  102. # GH#19959
  103. rng = timedelta_range('1 days 09:00:00', freq='H', periods=3)
  104. other = box([4, 3, 2])
  105. expected = TimedeltaIndex(['1 day 13:00:00'] * 3)
  106. with tm.assert_produces_warning(FutureWarning, check_stacklevel=False):
  107. # GH#22535
  108. result = rng + other
  109. tm.assert_index_equal(result, expected)
  110. with tm.assert_produces_warning(FutureWarning, check_stacklevel=False):
  111. # GH#22535
  112. result = other + rng
  113. tm.assert_index_equal(result, expected)
  114. @pytest.mark.parametrize('box', [np.array, pd.Index])
  115. def test_tdi_sub_integer_array(self, box):
  116. # GH#19959
  117. rng = timedelta_range('9H', freq='H', periods=3)
  118. other = box([4, 3, 2])
  119. expected = TimedeltaIndex(['5H', '7H', '9H'])
  120. with tm.assert_produces_warning(FutureWarning, check_stacklevel=False):
  121. # GH#22535
  122. result = rng - other
  123. tm.assert_index_equal(result, expected)
  124. with tm.assert_produces_warning(FutureWarning, check_stacklevel=False):
  125. # GH#22535
  126. result = other - rng
  127. tm.assert_index_equal(result, -expected)
  128. @pytest.mark.parametrize('box', [np.array, pd.Index])
  129. def test_tdi_addsub_integer_array_no_freq(self, box):
  130. # GH#19959
  131. tdi = TimedeltaIndex(['1 Day', 'NaT', '3 Hours'])
  132. other = box([14, -1, 16])
  133. with pytest.raises(NullFrequencyError):
  134. tdi + other
  135. with pytest.raises(NullFrequencyError):
  136. other + tdi
  137. with pytest.raises(NullFrequencyError):
  138. tdi - other
  139. with pytest.raises(NullFrequencyError):
  140. other - tdi
  141. # -------------------------------------------------------------
  142. # Binary operations TimedeltaIndex and timedelta-like
  143. # Note: add and sub are tested in tests.test_arithmetic, in-place
  144. # tests are kept here because their behavior is Index-specific
  145. def test_tdi_iadd_timedeltalike(self, delta):
  146. # only test adding/sub offsets as + is now numeric
  147. rng = timedelta_range('1 days', '10 days')
  148. expected = timedelta_range('1 days 02:00:00', '10 days 02:00:00',
  149. freq='D')
  150. rng += delta
  151. tm.assert_index_equal(rng, expected)
  152. def test_tdi_isub_timedeltalike(self, delta):
  153. # only test adding/sub offsets as - is now numeric
  154. rng = timedelta_range('1 days', '10 days')
  155. expected = timedelta_range('0 days 22:00:00', '9 days 22:00:00')
  156. rng -= delta
  157. tm.assert_index_equal(rng, expected)
  158. # -------------------------------------------------------------
  159. # TODO: after #24365 this probably belongs in scalar tests
  160. def test_ops_ndarray(self):
  161. td = Timedelta('1 day')
  162. # timedelta, timedelta
  163. other = pd.to_timedelta(['1 day']).values
  164. expected = pd.to_timedelta(['2 days']).values
  165. tm.assert_numpy_array_equal(td + other, expected)
  166. tm.assert_numpy_array_equal(other + td, expected)
  167. pytest.raises(TypeError, lambda: td + np.array([1]))
  168. pytest.raises(TypeError, lambda: np.array([1]) + td)
  169. expected = pd.to_timedelta(['0 days']).values
  170. tm.assert_numpy_array_equal(td - other, expected)
  171. tm.assert_numpy_array_equal(-other + td, expected)
  172. pytest.raises(TypeError, lambda: td - np.array([1]))
  173. pytest.raises(TypeError, lambda: np.array([1]) - td)
  174. expected = pd.to_timedelta(['2 days']).values
  175. tm.assert_numpy_array_equal(td * np.array([2]), expected)
  176. tm.assert_numpy_array_equal(np.array([2]) * td, expected)
  177. pytest.raises(TypeError, lambda: td * other)
  178. pytest.raises(TypeError, lambda: other * td)
  179. tm.assert_numpy_array_equal(td / other,
  180. np.array([1], dtype=np.float64))
  181. tm.assert_numpy_array_equal(other / td,
  182. np.array([1], dtype=np.float64))
  183. # timedelta, datetime
  184. other = pd.to_datetime(['2000-01-01']).values
  185. expected = pd.to_datetime(['2000-01-02']).values
  186. tm.assert_numpy_array_equal(td + other, expected)
  187. tm.assert_numpy_array_equal(other + td, expected)
  188. expected = pd.to_datetime(['1999-12-31']).values
  189. tm.assert_numpy_array_equal(-td + other, expected)
  190. tm.assert_numpy_array_equal(other - td, expected)
  191. def test_tdi_ops_attributes(self):
  192. rng = timedelta_range('2 days', periods=5, freq='2D', name='x')
  193. result = rng + 1 * rng.freq
  194. exp = timedelta_range('4 days', periods=5, freq='2D', name='x')
  195. tm.assert_index_equal(result, exp)
  196. assert result.freq == '2D'
  197. result = rng - 2 * rng.freq
  198. exp = timedelta_range('-2 days', periods=5, freq='2D', name='x')
  199. tm.assert_index_equal(result, exp)
  200. assert result.freq == '2D'
  201. result = rng * 2
  202. exp = timedelta_range('4 days', periods=5, freq='4D', name='x')
  203. tm.assert_index_equal(result, exp)
  204. assert result.freq == '4D'
  205. result = rng / 2
  206. exp = timedelta_range('1 days', periods=5, freq='D', name='x')
  207. tm.assert_index_equal(result, exp)
  208. assert result.freq == 'D'
  209. result = -rng
  210. exp = timedelta_range('-2 days', periods=5, freq='-2D', name='x')
  211. tm.assert_index_equal(result, exp)
  212. assert result.freq == '-2D'
  213. rng = pd.timedelta_range('-2 days', periods=5, freq='D', name='x')
  214. result = abs(rng)
  215. exp = TimedeltaIndex(['2 days', '1 days', '0 days', '1 days',
  216. '2 days'], name='x')
  217. tm.assert_index_equal(result, exp)
  218. assert result.freq is None