146 lines
4.6 KiB
Python
146 lines
4.6 KiB
Python
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)
|