test_insert.py 4.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145
  1. from datetime import timedelta
  2. import numpy as np
  3. import pytest
  4. from pandas._libs import lib
  5. import pandas as pd
  6. from pandas import (
  7. Index,
  8. Timedelta,
  9. TimedeltaIndex,
  10. timedelta_range,
  11. )
  12. import pandas._testing as tm
  13. class TestTimedeltaIndexInsert:
  14. def test_insert(self):
  15. idx = TimedeltaIndex(["4day", "1day", "2day"], name="idx")
  16. result = idx.insert(2, timedelta(days=5))
  17. exp = TimedeltaIndex(["4day", "1day", "5day", "2day"], name="idx")
  18. tm.assert_index_equal(result, exp)
  19. # insertion of non-datetime should coerce to object index
  20. result = idx.insert(1, "inserted")
  21. expected = Index(
  22. [Timedelta("4day"), "inserted", Timedelta("1day"), Timedelta("2day")],
  23. name="idx",
  24. )
  25. assert not isinstance(result, TimedeltaIndex)
  26. tm.assert_index_equal(result, expected)
  27. assert result.name == expected.name
  28. idx = timedelta_range("1day 00:00:01", periods=3, freq="s", name="idx")
  29. # preserve freq
  30. expected_0 = TimedeltaIndex(
  31. ["1day", "1day 00:00:01", "1day 00:00:02", "1day 00:00:03"],
  32. name="idx",
  33. freq="s",
  34. )
  35. expected_3 = TimedeltaIndex(
  36. ["1day 00:00:01", "1day 00:00:02", "1day 00:00:03", "1day 00:00:04"],
  37. name="idx",
  38. freq="s",
  39. )
  40. # reset freq to None
  41. expected_1_nofreq = TimedeltaIndex(
  42. ["1day 00:00:01", "1day 00:00:01", "1day 00:00:02", "1day 00:00:03"],
  43. name="idx",
  44. freq=None,
  45. )
  46. expected_3_nofreq = TimedeltaIndex(
  47. ["1day 00:00:01", "1day 00:00:02", "1day 00:00:03", "1day 00:00:05"],
  48. name="idx",
  49. freq=None,
  50. )
  51. cases = [
  52. (0, Timedelta("1day"), expected_0),
  53. (-3, Timedelta("1day"), expected_0),
  54. (3, Timedelta("1day 00:00:04"), expected_3),
  55. (1, Timedelta("1day 00:00:01"), expected_1_nofreq),
  56. (3, Timedelta("1day 00:00:05"), expected_3_nofreq),
  57. ]
  58. for n, d, expected in cases:
  59. result = idx.insert(n, d)
  60. tm.assert_index_equal(result, expected)
  61. assert result.name == expected.name
  62. assert result.freq == expected.freq
  63. @pytest.mark.parametrize(
  64. "null", [None, np.nan, np.timedelta64("NaT"), pd.NaT, pd.NA]
  65. )
  66. def test_insert_nat(self, null):
  67. # GH 18295 (test missing)
  68. idx = timedelta_range("1day", "3day")
  69. result = idx.insert(1, null)
  70. expected = TimedeltaIndex(["1day", pd.NaT, "2day", "3day"])
  71. tm.assert_index_equal(result, expected)
  72. def test_insert_invalid_na(self):
  73. idx = TimedeltaIndex(["4day", "1day", "2day"], name="idx")
  74. item = np.datetime64("NaT")
  75. result = idx.insert(0, item)
  76. expected = Index([item] + list(idx), dtype=object, name="idx")
  77. tm.assert_index_equal(result, expected)
  78. # Also works if we pass a different dt64nat object
  79. item2 = np.datetime64("NaT")
  80. result = idx.insert(0, item2)
  81. tm.assert_index_equal(result, expected)
  82. @pytest.mark.parametrize(
  83. "item", [0, np.int64(0), np.float64(0), np.array(0), np.datetime64(456, "us")]
  84. )
  85. def test_insert_mismatched_types_raises(self, item):
  86. # GH#33703 dont cast these to td64
  87. tdi = TimedeltaIndex(["4day", "1day", "2day"], name="idx")
  88. result = tdi.insert(1, item)
  89. expected = Index(
  90. [tdi[0], lib.item_from_zerodim(item)] + list(tdi[1:]),
  91. dtype=object,
  92. name="idx",
  93. )
  94. tm.assert_index_equal(result, expected)
  95. def test_insert_castable_str(self):
  96. idx = timedelta_range("1day", "3day")
  97. result = idx.insert(0, "1 Day")
  98. expected = TimedeltaIndex([idx[0]] + list(idx))
  99. tm.assert_index_equal(result, expected)
  100. def test_insert_non_castable_str(self):
  101. idx = timedelta_range("1day", "3day")
  102. result = idx.insert(0, "foo")
  103. expected = Index(["foo"] + list(idx), dtype=object)
  104. tm.assert_index_equal(result, expected)
  105. def test_insert_empty(self):
  106. # Corner case inserting with length zero doesn't raise IndexError
  107. # GH#33573 for freq preservation
  108. idx = timedelta_range("1 Day", periods=3)
  109. td = idx[0]
  110. result = idx[:0].insert(0, td)
  111. assert result.freq == "D"
  112. with pytest.raises(IndexError, match="loc must be an integer between"):
  113. result = idx[:0].insert(1, td)
  114. with pytest.raises(IndexError, match="loc must be an integer between"):
  115. result = idx[:0].insert(-1, td)