test_period.py 5.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182
  1. import numpy as np
  2. import pytest
  3. from pandas._libs.tslibs import iNaT
  4. from pandas._libs.tslibs.period import IncompatibleFrequency
  5. from pandas.core.dtypes.base import _registry as registry
  6. from pandas.core.dtypes.dtypes import PeriodDtype
  7. import pandas as pd
  8. import pandas._testing as tm
  9. from pandas.core.arrays import (
  10. PeriodArray,
  11. period_array,
  12. )
  13. # ----------------------------------------------------------------------------
  14. # Dtype
  15. def test_registered():
  16. assert PeriodDtype in registry.dtypes
  17. result = registry.find("Period[D]")
  18. expected = PeriodDtype("D")
  19. assert result == expected
  20. # ----------------------------------------------------------------------------
  21. # period_array
  22. def test_asi8():
  23. result = period_array(["2000", "2001", None], freq="D").asi8
  24. expected = np.array([10957, 11323, iNaT])
  25. tm.assert_numpy_array_equal(result, expected)
  26. def test_take_raises():
  27. arr = period_array(["2000", "2001"], freq="D")
  28. with pytest.raises(IncompatibleFrequency, match="freq"):
  29. arr.take([0, -1], allow_fill=True, fill_value=pd.Period("2000", freq="W"))
  30. msg = "value should be a 'Period' or 'NaT'. Got 'str' instead"
  31. with pytest.raises(TypeError, match=msg):
  32. arr.take([0, -1], allow_fill=True, fill_value="foo")
  33. def test_fillna_raises():
  34. arr = period_array(["2000", "2001", "2002"], freq="D")
  35. with pytest.raises(ValueError, match="Length"):
  36. arr.fillna(arr[:2])
  37. def test_fillna_copies():
  38. arr = period_array(["2000", "2001", "2002"], freq="D")
  39. result = arr.fillna(pd.Period("2000", "D"))
  40. assert result is not arr
  41. # ----------------------------------------------------------------------------
  42. # setitem
  43. @pytest.mark.parametrize(
  44. "key, value, expected",
  45. [
  46. ([0], pd.Period("2000", "D"), [10957, 1, 2]),
  47. ([0], None, [iNaT, 1, 2]),
  48. ([0], np.nan, [iNaT, 1, 2]),
  49. ([0, 1, 2], pd.Period("2000", "D"), [10957] * 3),
  50. (
  51. [0, 1, 2],
  52. [pd.Period("2000", "D"), pd.Period("2001", "D"), pd.Period("2002", "D")],
  53. [10957, 11323, 11688],
  54. ),
  55. ],
  56. )
  57. def test_setitem(key, value, expected):
  58. arr = PeriodArray(np.arange(3), freq="D")
  59. expected = PeriodArray(expected, freq="D")
  60. arr[key] = value
  61. tm.assert_period_array_equal(arr, expected)
  62. def test_setitem_raises_incompatible_freq():
  63. arr = PeriodArray(np.arange(3), freq="D")
  64. with pytest.raises(IncompatibleFrequency, match="freq"):
  65. arr[0] = pd.Period("2000", freq="A")
  66. other = period_array(["2000", "2001"], freq="A")
  67. with pytest.raises(IncompatibleFrequency, match="freq"):
  68. arr[[0, 1]] = other
  69. def test_setitem_raises_length():
  70. arr = PeriodArray(np.arange(3), freq="D")
  71. with pytest.raises(ValueError, match="length"):
  72. arr[[0, 1]] = [pd.Period("2000", freq="D")]
  73. def test_setitem_raises_type():
  74. arr = PeriodArray(np.arange(3), freq="D")
  75. with pytest.raises(TypeError, match="int"):
  76. arr[0] = 1
  77. # ----------------------------------------------------------------------------
  78. # Ops
  79. def test_sub_period():
  80. arr = period_array(["2000", "2001"], freq="D")
  81. other = pd.Period("2000", freq="M")
  82. with pytest.raises(IncompatibleFrequency, match="freq"):
  83. arr - other
  84. def test_sub_period_overflow():
  85. # GH#47538
  86. dti = pd.date_range("1677-09-22", periods=2, freq="D")
  87. pi = dti.to_period("ns")
  88. per = pd.Period._from_ordinal(10**14, pi.freq)
  89. with pytest.raises(OverflowError, match="Overflow in int64 addition"):
  90. pi - per
  91. with pytest.raises(OverflowError, match="Overflow in int64 addition"):
  92. per - pi
  93. # ----------------------------------------------------------------------------
  94. # Methods
  95. @pytest.mark.parametrize(
  96. "other",
  97. [pd.Period("2000", freq="H"), period_array(["2000", "2001", "2000"], freq="H")],
  98. )
  99. def test_where_different_freq_raises(other):
  100. # GH#45768 The PeriodArray method raises, the Series method coerces
  101. ser = pd.Series(period_array(["2000", "2001", "2002"], freq="D"))
  102. cond = np.array([True, False, True])
  103. with pytest.raises(IncompatibleFrequency, match="freq"):
  104. ser.array._where(cond, other)
  105. res = ser.where(cond, other)
  106. expected = ser.astype(object).where(cond, other)
  107. tm.assert_series_equal(res, expected)
  108. # ----------------------------------------------------------------------------
  109. # Printing
  110. def test_repr_small():
  111. arr = period_array(["2000", "2001"], freq="D")
  112. result = str(arr)
  113. expected = (
  114. "<PeriodArray>\n['2000-01-01', '2001-01-01']\nLength: 2, dtype: period[D]"
  115. )
  116. assert result == expected
  117. def test_repr_large():
  118. arr = period_array(["2000", "2001"] * 500, freq="D")
  119. result = str(arr)
  120. expected = (
  121. "<PeriodArray>\n"
  122. "['2000-01-01', '2001-01-01', '2000-01-01', '2001-01-01', "
  123. "'2000-01-01',\n"
  124. " '2001-01-01', '2000-01-01', '2001-01-01', '2000-01-01', "
  125. "'2001-01-01',\n"
  126. " ...\n"
  127. " '2000-01-01', '2001-01-01', '2000-01-01', '2001-01-01', "
  128. "'2000-01-01',\n"
  129. " '2001-01-01', '2000-01-01', '2001-01-01', '2000-01-01', "
  130. "'2001-01-01']\n"
  131. "Length: 1000, dtype: period[D]"
  132. )
  133. assert result == expected