test_delete.py 4.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138
  1. import pytest
  2. from pandas import (
  3. DatetimeIndex,
  4. Series,
  5. date_range,
  6. )
  7. import pandas._testing as tm
  8. class TestDelete:
  9. def test_delete(self):
  10. idx = date_range(start="2000-01-01", periods=5, freq="M", name="idx")
  11. # preserve freq
  12. expected_0 = date_range(start="2000-02-01", periods=4, freq="M", name="idx")
  13. expected_4 = date_range(start="2000-01-01", periods=4, freq="M", name="idx")
  14. # reset freq to None
  15. expected_1 = DatetimeIndex(
  16. ["2000-01-31", "2000-03-31", "2000-04-30", "2000-05-31"],
  17. freq=None,
  18. name="idx",
  19. )
  20. cases = {
  21. 0: expected_0,
  22. -5: expected_0,
  23. -1: expected_4,
  24. 4: expected_4,
  25. 1: expected_1,
  26. }
  27. for n, expected in cases.items():
  28. result = idx.delete(n)
  29. tm.assert_index_equal(result, expected)
  30. assert result.name == expected.name
  31. assert result.freq == expected.freq
  32. with pytest.raises((IndexError, ValueError), match="out of bounds"):
  33. # either depending on numpy version
  34. idx.delete(5)
  35. for tz in [None, "Asia/Tokyo", "US/Pacific"]:
  36. idx = date_range(
  37. start="2000-01-01 09:00", periods=10, freq="H", name="idx", tz=tz
  38. )
  39. expected = date_range(
  40. start="2000-01-01 10:00", periods=9, freq="H", name="idx", tz=tz
  41. )
  42. result = idx.delete(0)
  43. tm.assert_index_equal(result, expected)
  44. assert result.name == expected.name
  45. assert result.freqstr == "H"
  46. assert result.tz == expected.tz
  47. expected = date_range(
  48. start="2000-01-01 09:00", periods=9, freq="H", name="idx", tz=tz
  49. )
  50. result = idx.delete(-1)
  51. tm.assert_index_equal(result, expected)
  52. assert result.name == expected.name
  53. assert result.freqstr == "H"
  54. assert result.tz == expected.tz
  55. def test_delete_slice(self):
  56. idx = date_range(start="2000-01-01", periods=10, freq="D", name="idx")
  57. # preserve freq
  58. expected_0_2 = date_range(start="2000-01-04", periods=7, freq="D", name="idx")
  59. expected_7_9 = date_range(start="2000-01-01", periods=7, freq="D", name="idx")
  60. # reset freq to None
  61. expected_3_5 = DatetimeIndex(
  62. [
  63. "2000-01-01",
  64. "2000-01-02",
  65. "2000-01-03",
  66. "2000-01-07",
  67. "2000-01-08",
  68. "2000-01-09",
  69. "2000-01-10",
  70. ],
  71. freq=None,
  72. name="idx",
  73. )
  74. cases = {
  75. (0, 1, 2): expected_0_2,
  76. (7, 8, 9): expected_7_9,
  77. (3, 4, 5): expected_3_5,
  78. }
  79. for n, expected in cases.items():
  80. result = idx.delete(n)
  81. tm.assert_index_equal(result, expected)
  82. assert result.name == expected.name
  83. assert result.freq == expected.freq
  84. result = idx.delete(slice(n[0], n[-1] + 1))
  85. tm.assert_index_equal(result, expected)
  86. assert result.name == expected.name
  87. assert result.freq == expected.freq
  88. for tz in [None, "Asia/Tokyo", "US/Pacific"]:
  89. ts = Series(
  90. 1,
  91. index=date_range(
  92. "2000-01-01 09:00", periods=10, freq="H", name="idx", tz=tz
  93. ),
  94. )
  95. # preserve freq
  96. result = ts.drop(ts.index[:5]).index
  97. expected = date_range(
  98. "2000-01-01 14:00", periods=5, freq="H", name="idx", tz=tz
  99. )
  100. tm.assert_index_equal(result, expected)
  101. assert result.name == expected.name
  102. assert result.freq == expected.freq
  103. assert result.tz == expected.tz
  104. # reset freq to None
  105. result = ts.drop(ts.index[[1, 3, 5, 7, 9]]).index
  106. expected = DatetimeIndex(
  107. [
  108. "2000-01-01 09:00",
  109. "2000-01-01 11:00",
  110. "2000-01-01 13:00",
  111. "2000-01-01 15:00",
  112. "2000-01-01 17:00",
  113. ],
  114. freq=None,
  115. name="idx",
  116. tz=tz,
  117. )
  118. tm.assert_index_equal(result, expected)
  119. assert result.name == expected.name
  120. assert result.freq == expected.freq
  121. assert result.tz == expected.tz