test_insert.py 8.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256
  1. from datetime import datetime
  2. import numpy as np
  3. import pytest
  4. import pytz
  5. from pandas import (
  6. NA,
  7. DatetimeIndex,
  8. Index,
  9. NaT,
  10. Timestamp,
  11. date_range,
  12. )
  13. import pandas._testing as tm
  14. class TestInsert:
  15. @pytest.mark.parametrize("null", [None, np.nan, np.datetime64("NaT"), NaT, NA])
  16. @pytest.mark.parametrize("tz", [None, "UTC", "US/Eastern"])
  17. def test_insert_nat(self, tz, null):
  18. # GH#16537, GH#18295 (test missing)
  19. idx = DatetimeIndex(["2017-01-01"], tz=tz)
  20. expected = DatetimeIndex(["NaT", "2017-01-01"], tz=tz)
  21. if tz is not None and isinstance(null, np.datetime64):
  22. expected = Index([null, idx[0]], dtype=object)
  23. res = idx.insert(0, null)
  24. tm.assert_index_equal(res, expected)
  25. @pytest.mark.parametrize("tz", [None, "UTC", "US/Eastern"])
  26. def test_insert_invalid_na(self, tz):
  27. idx = DatetimeIndex(["2017-01-01"], tz=tz)
  28. item = np.timedelta64("NaT")
  29. result = idx.insert(0, item)
  30. expected = Index([item] + list(idx), dtype=object)
  31. tm.assert_index_equal(result, expected)
  32. def test_insert_empty_preserves_freq(self, tz_naive_fixture):
  33. # GH#33573
  34. tz = tz_naive_fixture
  35. dti = DatetimeIndex([], tz=tz, freq="D")
  36. item = Timestamp("2017-04-05").tz_localize(tz)
  37. result = dti.insert(0, item)
  38. assert result.freq == dti.freq
  39. # But not when we insert an item that doesn't conform to freq
  40. dti = DatetimeIndex([], tz=tz, freq="W-THU")
  41. result = dti.insert(0, item)
  42. assert result.freq is None
  43. def test_insert(self):
  44. idx = DatetimeIndex(["2000-01-04", "2000-01-01", "2000-01-02"], name="idx")
  45. result = idx.insert(2, datetime(2000, 1, 5))
  46. exp = DatetimeIndex(
  47. ["2000-01-04", "2000-01-01", "2000-01-05", "2000-01-02"], name="idx"
  48. )
  49. tm.assert_index_equal(result, exp)
  50. # insertion of non-datetime should coerce to object index
  51. result = idx.insert(1, "inserted")
  52. expected = Index(
  53. [
  54. datetime(2000, 1, 4),
  55. "inserted",
  56. datetime(2000, 1, 1),
  57. datetime(2000, 1, 2),
  58. ],
  59. name="idx",
  60. )
  61. assert not isinstance(result, DatetimeIndex)
  62. tm.assert_index_equal(result, expected)
  63. assert result.name == expected.name
  64. idx = date_range("1/1/2000", periods=3, freq="M", name="idx")
  65. # preserve freq
  66. expected_0 = DatetimeIndex(
  67. ["1999-12-31", "2000-01-31", "2000-02-29", "2000-03-31"],
  68. name="idx",
  69. freq="M",
  70. )
  71. expected_3 = DatetimeIndex(
  72. ["2000-01-31", "2000-02-29", "2000-03-31", "2000-04-30"],
  73. name="idx",
  74. freq="M",
  75. )
  76. # reset freq to None
  77. expected_1_nofreq = DatetimeIndex(
  78. ["2000-01-31", "2000-01-31", "2000-02-29", "2000-03-31"],
  79. name="idx",
  80. freq=None,
  81. )
  82. expected_3_nofreq = DatetimeIndex(
  83. ["2000-01-31", "2000-02-29", "2000-03-31", "2000-01-02"],
  84. name="idx",
  85. freq=None,
  86. )
  87. cases = [
  88. (0, datetime(1999, 12, 31), expected_0),
  89. (-3, datetime(1999, 12, 31), expected_0),
  90. (3, datetime(2000, 4, 30), expected_3),
  91. (1, datetime(2000, 1, 31), expected_1_nofreq),
  92. (3, datetime(2000, 1, 2), expected_3_nofreq),
  93. ]
  94. for n, d, expected in cases:
  95. result = idx.insert(n, d)
  96. tm.assert_index_equal(result, expected)
  97. assert result.name == expected.name
  98. assert result.freq == expected.freq
  99. # reset freq to None
  100. result = idx.insert(3, datetime(2000, 1, 2))
  101. expected = DatetimeIndex(
  102. ["2000-01-31", "2000-02-29", "2000-03-31", "2000-01-02"],
  103. name="idx",
  104. freq=None,
  105. )
  106. tm.assert_index_equal(result, expected)
  107. assert result.name == expected.name
  108. assert result.freq is None
  109. for tz in ["US/Pacific", "Asia/Singapore"]:
  110. idx = date_range("1/1/2000 09:00", periods=6, freq="H", tz=tz, name="idx")
  111. # preserve freq
  112. expected = date_range(
  113. "1/1/2000 09:00", periods=7, freq="H", tz=tz, name="idx"
  114. )
  115. for d in [
  116. Timestamp("2000-01-01 15:00", tz=tz),
  117. pytz.timezone(tz).localize(datetime(2000, 1, 1, 15)),
  118. ]:
  119. result = idx.insert(6, d)
  120. tm.assert_index_equal(result, expected)
  121. assert result.name == expected.name
  122. assert result.freq == expected.freq
  123. assert result.tz == expected.tz
  124. expected = DatetimeIndex(
  125. [
  126. "2000-01-01 09:00",
  127. "2000-01-01 10:00",
  128. "2000-01-01 11:00",
  129. "2000-01-01 12:00",
  130. "2000-01-01 13:00",
  131. "2000-01-01 14:00",
  132. "2000-01-01 10:00",
  133. ],
  134. name="idx",
  135. tz=tz,
  136. freq=None,
  137. )
  138. # reset freq to None
  139. for d in [
  140. Timestamp("2000-01-01 10:00", tz=tz),
  141. pytz.timezone(tz).localize(datetime(2000, 1, 1, 10)),
  142. ]:
  143. result = idx.insert(6, d)
  144. tm.assert_index_equal(result, expected)
  145. assert result.name == expected.name
  146. assert result.tz == expected.tz
  147. assert result.freq is None
  148. # TODO: also changes DataFrame.__setitem__ with expansion
  149. def test_insert_mismatched_tzawareness(self):
  150. # see GH#7299
  151. idx = date_range("1/1/2000", periods=3, freq="D", tz="Asia/Tokyo", name="idx")
  152. # mismatched tz-awareness
  153. item = Timestamp("2000-01-04")
  154. result = idx.insert(3, item)
  155. expected = Index(
  156. list(idx[:3]) + [item] + list(idx[3:]), dtype=object, name="idx"
  157. )
  158. tm.assert_index_equal(result, expected)
  159. # mismatched tz-awareness
  160. item = datetime(2000, 1, 4)
  161. result = idx.insert(3, item)
  162. expected = Index(
  163. list(idx[:3]) + [item] + list(idx[3:]), dtype=object, name="idx"
  164. )
  165. tm.assert_index_equal(result, expected)
  166. # TODO: also changes DataFrame.__setitem__ with expansion
  167. def test_insert_mismatched_tz(self):
  168. # see GH#7299
  169. # pre-2.0 with mismatched tzs we would cast to object
  170. idx = date_range("1/1/2000", periods=3, freq="D", tz="Asia/Tokyo", name="idx")
  171. # mismatched tz -> cast to object (could reasonably cast to same tz or UTC)
  172. item = Timestamp("2000-01-04", tz="US/Eastern")
  173. result = idx.insert(3, item)
  174. expected = Index(
  175. list(idx[:3]) + [item.tz_convert(idx.tz)] + list(idx[3:]),
  176. name="idx",
  177. )
  178. assert expected.dtype == idx.dtype
  179. tm.assert_index_equal(result, expected)
  180. item = datetime(2000, 1, 4, tzinfo=pytz.timezone("US/Eastern"))
  181. result = idx.insert(3, item)
  182. expected = Index(
  183. list(idx[:3]) + [item.astimezone(idx.tzinfo)] + list(idx[3:]),
  184. name="idx",
  185. )
  186. assert expected.dtype == idx.dtype
  187. tm.assert_index_equal(result, expected)
  188. @pytest.mark.parametrize(
  189. "item", [0, np.int64(0), np.float64(0), np.array(0), np.timedelta64(456)]
  190. )
  191. def test_insert_mismatched_types_raises(self, tz_aware_fixture, item):
  192. # GH#33703 dont cast these to dt64
  193. tz = tz_aware_fixture
  194. dti = date_range("2019-11-04", periods=9, freq="-1D", name=9, tz=tz)
  195. result = dti.insert(1, item)
  196. if isinstance(item, np.ndarray):
  197. assert item.item() == 0
  198. expected = Index([dti[0], 0] + list(dti[1:]), dtype=object, name=9)
  199. else:
  200. expected = Index([dti[0], item] + list(dti[1:]), dtype=object, name=9)
  201. tm.assert_index_equal(result, expected)
  202. def test_insert_castable_str(self, tz_aware_fixture):
  203. # GH#33703
  204. tz = tz_aware_fixture
  205. dti = date_range("2019-11-04", periods=3, freq="-1D", name=9, tz=tz)
  206. value = "2019-11-05"
  207. result = dti.insert(0, value)
  208. ts = Timestamp(value).tz_localize(tz)
  209. expected = DatetimeIndex([ts] + list(dti), dtype=dti.dtype, name=9)
  210. tm.assert_index_equal(result, expected)
  211. def test_insert_non_castable_str(self, tz_aware_fixture):
  212. # GH#33703
  213. tz = tz_aware_fixture
  214. dti = date_range("2019-11-04", periods=3, freq="-1D", name=9, tz=tz)
  215. value = "foo"
  216. result = dti.insert(0, value)
  217. expected = Index(["foo"] + list(dti), dtype=object, name=9)
  218. tm.assert_index_equal(result, expected)