test_year.py 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339
  1. """
  2. Tests for the following offsets:
  3. - YearBegin
  4. - YearEnd
  5. """
  6. from __future__ import annotations
  7. from datetime import datetime
  8. import numpy as np
  9. import pytest
  10. from pandas import Timestamp
  11. from pandas.tests.tseries.offsets.common import (
  12. assert_is_on_offset,
  13. assert_offset_equal,
  14. )
  15. from pandas.tseries.offsets import (
  16. YearBegin,
  17. YearEnd,
  18. )
  19. class TestYearBegin:
  20. def test_misspecified(self):
  21. with pytest.raises(ValueError, match="Month must go from 1 to 12"):
  22. YearBegin(month=13)
  23. offset_cases = []
  24. offset_cases.append(
  25. (
  26. YearBegin(),
  27. {
  28. datetime(2008, 1, 1): datetime(2009, 1, 1),
  29. datetime(2008, 6, 30): datetime(2009, 1, 1),
  30. datetime(2008, 12, 31): datetime(2009, 1, 1),
  31. datetime(2005, 12, 30): datetime(2006, 1, 1),
  32. datetime(2005, 12, 31): datetime(2006, 1, 1),
  33. },
  34. )
  35. )
  36. offset_cases.append(
  37. (
  38. YearBegin(0),
  39. {
  40. datetime(2008, 1, 1): datetime(2008, 1, 1),
  41. datetime(2008, 6, 30): datetime(2009, 1, 1),
  42. datetime(2008, 12, 31): datetime(2009, 1, 1),
  43. datetime(2005, 12, 30): datetime(2006, 1, 1),
  44. datetime(2005, 12, 31): datetime(2006, 1, 1),
  45. },
  46. )
  47. )
  48. offset_cases.append(
  49. (
  50. YearBegin(3),
  51. {
  52. datetime(2008, 1, 1): datetime(2011, 1, 1),
  53. datetime(2008, 6, 30): datetime(2011, 1, 1),
  54. datetime(2008, 12, 31): datetime(2011, 1, 1),
  55. datetime(2005, 12, 30): datetime(2008, 1, 1),
  56. datetime(2005, 12, 31): datetime(2008, 1, 1),
  57. },
  58. )
  59. )
  60. offset_cases.append(
  61. (
  62. YearBegin(-1),
  63. {
  64. datetime(2007, 1, 1): datetime(2006, 1, 1),
  65. datetime(2007, 1, 15): datetime(2007, 1, 1),
  66. datetime(2008, 6, 30): datetime(2008, 1, 1),
  67. datetime(2008, 12, 31): datetime(2008, 1, 1),
  68. datetime(2006, 12, 29): datetime(2006, 1, 1),
  69. datetime(2006, 12, 30): datetime(2006, 1, 1),
  70. datetime(2007, 1, 1): datetime(2006, 1, 1),
  71. },
  72. )
  73. )
  74. offset_cases.append(
  75. (
  76. YearBegin(-2),
  77. {
  78. datetime(2007, 1, 1): datetime(2005, 1, 1),
  79. datetime(2008, 6, 30): datetime(2007, 1, 1),
  80. datetime(2008, 12, 31): datetime(2007, 1, 1),
  81. },
  82. )
  83. )
  84. offset_cases.append(
  85. (
  86. YearBegin(month=4),
  87. {
  88. datetime(2007, 4, 1): datetime(2008, 4, 1),
  89. datetime(2007, 4, 15): datetime(2008, 4, 1),
  90. datetime(2007, 3, 1): datetime(2007, 4, 1),
  91. datetime(2007, 12, 15): datetime(2008, 4, 1),
  92. datetime(2012, 1, 31): datetime(2012, 4, 1),
  93. },
  94. )
  95. )
  96. offset_cases.append(
  97. (
  98. YearBegin(0, month=4),
  99. {
  100. datetime(2007, 4, 1): datetime(2007, 4, 1),
  101. datetime(2007, 3, 1): datetime(2007, 4, 1),
  102. datetime(2007, 12, 15): datetime(2008, 4, 1),
  103. datetime(2012, 1, 31): datetime(2012, 4, 1),
  104. },
  105. )
  106. )
  107. offset_cases.append(
  108. (
  109. YearBegin(4, month=4),
  110. {
  111. datetime(2007, 4, 1): datetime(2011, 4, 1),
  112. datetime(2007, 4, 15): datetime(2011, 4, 1),
  113. datetime(2007, 3, 1): datetime(2010, 4, 1),
  114. datetime(2007, 12, 15): datetime(2011, 4, 1),
  115. datetime(2012, 1, 31): datetime(2015, 4, 1),
  116. },
  117. )
  118. )
  119. offset_cases.append(
  120. (
  121. YearBegin(-1, month=4),
  122. {
  123. datetime(2007, 4, 1): datetime(2006, 4, 1),
  124. datetime(2007, 3, 1): datetime(2006, 4, 1),
  125. datetime(2007, 12, 15): datetime(2007, 4, 1),
  126. datetime(2012, 1, 31): datetime(2011, 4, 1),
  127. },
  128. )
  129. )
  130. offset_cases.append(
  131. (
  132. YearBegin(-3, month=4),
  133. {
  134. datetime(2007, 4, 1): datetime(2004, 4, 1),
  135. datetime(2007, 3, 1): datetime(2004, 4, 1),
  136. datetime(2007, 12, 15): datetime(2005, 4, 1),
  137. datetime(2012, 1, 31): datetime(2009, 4, 1),
  138. },
  139. )
  140. )
  141. @pytest.mark.parametrize("case", offset_cases)
  142. def test_offset(self, case):
  143. offset, cases = case
  144. for base, expected in cases.items():
  145. assert_offset_equal(offset, base, expected)
  146. on_offset_cases = [
  147. (YearBegin(), datetime(2007, 1, 3), False),
  148. (YearBegin(), datetime(2008, 1, 1), True),
  149. (YearBegin(), datetime(2006, 12, 31), False),
  150. (YearBegin(), datetime(2006, 1, 2), False),
  151. ]
  152. @pytest.mark.parametrize("case", on_offset_cases)
  153. def test_is_on_offset(self, case):
  154. offset, dt, expected = case
  155. assert_is_on_offset(offset, dt, expected)
  156. class TestYearEnd:
  157. def test_misspecified(self):
  158. with pytest.raises(ValueError, match="Month must go from 1 to 12"):
  159. YearEnd(month=13)
  160. offset_cases = []
  161. offset_cases.append(
  162. (
  163. YearEnd(),
  164. {
  165. datetime(2008, 1, 1): datetime(2008, 12, 31),
  166. datetime(2008, 6, 30): datetime(2008, 12, 31),
  167. datetime(2008, 12, 31): datetime(2009, 12, 31),
  168. datetime(2005, 12, 30): datetime(2005, 12, 31),
  169. datetime(2005, 12, 31): datetime(2006, 12, 31),
  170. },
  171. )
  172. )
  173. offset_cases.append(
  174. (
  175. YearEnd(0),
  176. {
  177. datetime(2008, 1, 1): datetime(2008, 12, 31),
  178. datetime(2008, 6, 30): datetime(2008, 12, 31),
  179. datetime(2008, 12, 31): datetime(2008, 12, 31),
  180. datetime(2005, 12, 30): datetime(2005, 12, 31),
  181. },
  182. )
  183. )
  184. offset_cases.append(
  185. (
  186. YearEnd(-1),
  187. {
  188. datetime(2007, 1, 1): datetime(2006, 12, 31),
  189. datetime(2008, 6, 30): datetime(2007, 12, 31),
  190. datetime(2008, 12, 31): datetime(2007, 12, 31),
  191. datetime(2006, 12, 29): datetime(2005, 12, 31),
  192. datetime(2006, 12, 30): datetime(2005, 12, 31),
  193. datetime(2007, 1, 1): datetime(2006, 12, 31),
  194. },
  195. )
  196. )
  197. offset_cases.append(
  198. (
  199. YearEnd(-2),
  200. {
  201. datetime(2007, 1, 1): datetime(2005, 12, 31),
  202. datetime(2008, 6, 30): datetime(2006, 12, 31),
  203. datetime(2008, 12, 31): datetime(2006, 12, 31),
  204. },
  205. )
  206. )
  207. @pytest.mark.parametrize("case", offset_cases)
  208. def test_offset(self, case):
  209. offset, cases = case
  210. for base, expected in cases.items():
  211. assert_offset_equal(offset, base, expected)
  212. on_offset_cases = [
  213. (YearEnd(), datetime(2007, 12, 31), True),
  214. (YearEnd(), datetime(2008, 1, 1), False),
  215. (YearEnd(), datetime(2006, 12, 31), True),
  216. (YearEnd(), datetime(2006, 12, 29), False),
  217. ]
  218. @pytest.mark.parametrize("case", on_offset_cases)
  219. def test_is_on_offset(self, case):
  220. offset, dt, expected = case
  221. assert_is_on_offset(offset, dt, expected)
  222. class TestYearEndDiffMonth:
  223. offset_cases = []
  224. offset_cases.append(
  225. (
  226. YearEnd(month=3),
  227. {
  228. datetime(2008, 1, 1): datetime(2008, 3, 31),
  229. datetime(2008, 2, 15): datetime(2008, 3, 31),
  230. datetime(2008, 3, 31): datetime(2009, 3, 31),
  231. datetime(2008, 3, 30): datetime(2008, 3, 31),
  232. datetime(2005, 3, 31): datetime(2006, 3, 31),
  233. datetime(2006, 7, 30): datetime(2007, 3, 31),
  234. },
  235. )
  236. )
  237. offset_cases.append(
  238. (
  239. YearEnd(0, month=3),
  240. {
  241. datetime(2008, 1, 1): datetime(2008, 3, 31),
  242. datetime(2008, 2, 28): datetime(2008, 3, 31),
  243. datetime(2008, 3, 31): datetime(2008, 3, 31),
  244. datetime(2005, 3, 30): datetime(2005, 3, 31),
  245. },
  246. )
  247. )
  248. offset_cases.append(
  249. (
  250. YearEnd(-1, month=3),
  251. {
  252. datetime(2007, 1, 1): datetime(2006, 3, 31),
  253. datetime(2008, 2, 28): datetime(2007, 3, 31),
  254. datetime(2008, 3, 31): datetime(2007, 3, 31),
  255. datetime(2006, 3, 29): datetime(2005, 3, 31),
  256. datetime(2006, 3, 30): datetime(2005, 3, 31),
  257. datetime(2007, 3, 1): datetime(2006, 3, 31),
  258. },
  259. )
  260. )
  261. offset_cases.append(
  262. (
  263. YearEnd(-2, month=3),
  264. {
  265. datetime(2007, 1, 1): datetime(2005, 3, 31),
  266. datetime(2008, 6, 30): datetime(2007, 3, 31),
  267. datetime(2008, 3, 31): datetime(2006, 3, 31),
  268. },
  269. )
  270. )
  271. @pytest.mark.parametrize("case", offset_cases)
  272. def test_offset(self, case):
  273. offset, cases = case
  274. for base, expected in cases.items():
  275. assert_offset_equal(offset, base, expected)
  276. on_offset_cases = [
  277. (YearEnd(month=3), datetime(2007, 3, 31), True),
  278. (YearEnd(month=3), datetime(2008, 1, 1), False),
  279. (YearEnd(month=3), datetime(2006, 3, 31), True),
  280. (YearEnd(month=3), datetime(2006, 3, 29), False),
  281. ]
  282. @pytest.mark.parametrize("case", on_offset_cases)
  283. def test_is_on_offset(self, case):
  284. offset, dt, expected = case
  285. assert_is_on_offset(offset, dt, expected)
  286. def test_add_out_of_pydatetime_range():
  287. # GH#50348 don't raise in Timestamp.replace
  288. ts = Timestamp(np.datetime64("-20000-12-31"))
  289. off = YearEnd()
  290. result = ts + off
  291. # TODO(cython3): "arg: datetime" annotation will impose
  292. # datetime limitations on Timestamp. The fused type below works in cy3
  293. # ctypedef fused datetimelike:
  294. # _Timestamp
  295. # datetime
  296. # expected = Timestamp(np.datetime64("-19999-12-31"))
  297. # assert result == expected
  298. assert result.year in (-19999, 1973)
  299. assert result.month == 12
  300. assert result.day == 31