123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145 |
- from datetime import timedelta
- import numpy as np
- import pytest
- from pandas._libs import lib
- import pandas as pd
- from pandas import (
- Index,
- Timedelta,
- TimedeltaIndex,
- timedelta_range,
- )
- import pandas._testing as tm
- class TestTimedeltaIndexInsert:
- def test_insert(self):
- idx = TimedeltaIndex(["4day", "1day", "2day"], name="idx")
- result = idx.insert(2, timedelta(days=5))
- exp = TimedeltaIndex(["4day", "1day", "5day", "2day"], name="idx")
- tm.assert_index_equal(result, exp)
- # insertion of non-datetime should coerce to object index
- result = idx.insert(1, "inserted")
- expected = Index(
- [Timedelta("4day"), "inserted", Timedelta("1day"), Timedelta("2day")],
- name="idx",
- )
- assert not isinstance(result, TimedeltaIndex)
- tm.assert_index_equal(result, expected)
- assert result.name == expected.name
- idx = timedelta_range("1day 00:00:01", periods=3, freq="s", name="idx")
- # preserve freq
- expected_0 = TimedeltaIndex(
- ["1day", "1day 00:00:01", "1day 00:00:02", "1day 00:00:03"],
- name="idx",
- freq="s",
- )
- expected_3 = TimedeltaIndex(
- ["1day 00:00:01", "1day 00:00:02", "1day 00:00:03", "1day 00:00:04"],
- name="idx",
- freq="s",
- )
- # reset freq to None
- expected_1_nofreq = TimedeltaIndex(
- ["1day 00:00:01", "1day 00:00:01", "1day 00:00:02", "1day 00:00:03"],
- name="idx",
- freq=None,
- )
- expected_3_nofreq = TimedeltaIndex(
- ["1day 00:00:01", "1day 00:00:02", "1day 00:00:03", "1day 00:00:05"],
- name="idx",
- freq=None,
- )
- cases = [
- (0, Timedelta("1day"), expected_0),
- (-3, Timedelta("1day"), expected_0),
- (3, Timedelta("1day 00:00:04"), expected_3),
- (1, Timedelta("1day 00:00:01"), expected_1_nofreq),
- (3, Timedelta("1day 00:00:05"), expected_3_nofreq),
- ]
- for n, d, expected in cases:
- result = idx.insert(n, d)
- tm.assert_index_equal(result, expected)
- assert result.name == expected.name
- assert result.freq == expected.freq
- @pytest.mark.parametrize(
- "null", [None, np.nan, np.timedelta64("NaT"), pd.NaT, pd.NA]
- )
- def test_insert_nat(self, null):
- # GH 18295 (test missing)
- idx = timedelta_range("1day", "3day")
- result = idx.insert(1, null)
- expected = TimedeltaIndex(["1day", pd.NaT, "2day", "3day"])
- tm.assert_index_equal(result, expected)
- def test_insert_invalid_na(self):
- idx = TimedeltaIndex(["4day", "1day", "2day"], name="idx")
- item = np.datetime64("NaT")
- result = idx.insert(0, item)
- expected = Index([item] + list(idx), dtype=object, name="idx")
- tm.assert_index_equal(result, expected)
- # Also works if we pass a different dt64nat object
- item2 = np.datetime64("NaT")
- result = idx.insert(0, item2)
- tm.assert_index_equal(result, expected)
- @pytest.mark.parametrize(
- "item", [0, np.int64(0), np.float64(0), np.array(0), np.datetime64(456, "us")]
- )
- def test_insert_mismatched_types_raises(self, item):
- # GH#33703 dont cast these to td64
- tdi = TimedeltaIndex(["4day", "1day", "2day"], name="idx")
- result = tdi.insert(1, item)
- expected = Index(
- [tdi[0], lib.item_from_zerodim(item)] + list(tdi[1:]),
- dtype=object,
- name="idx",
- )
- tm.assert_index_equal(result, expected)
- def test_insert_castable_str(self):
- idx = timedelta_range("1day", "3day")
- result = idx.insert(0, "1 Day")
- expected = TimedeltaIndex([idx[0]] + list(idx))
- tm.assert_index_equal(result, expected)
- def test_insert_non_castable_str(self):
- idx = timedelta_range("1day", "3day")
- result = idx.insert(0, "foo")
- expected = Index(["foo"] + list(idx), dtype=object)
- tm.assert_index_equal(result, expected)
- def test_insert_empty(self):
- # Corner case inserting with length zero doesn't raise IndexError
- # GH#33573 for freq preservation
- idx = timedelta_range("1 Day", periods=3)
- td = idx[0]
- result = idx[:0].insert(0, td)
- assert result.freq == "D"
- with pytest.raises(IndexError, match="loc must be an integer between"):
- result = idx[:0].insert(1, td)
- with pytest.raises(IndexError, match="loc must be an integer between"):
- result = idx[:0].insert(-1, td)
|