85 lines
2.9 KiB
Python
85 lines
2.9 KiB
Python
import numpy as np
|
|
import pytest
|
|
|
|
from pandas._libs.tslibs import IncompatibleFrequency
|
|
from pandas.compat.numpy import np_version_under1p18
|
|
|
|
from pandas import NaT, Period, PeriodIndex, Series, array
|
|
import pandas._testing as tm
|
|
|
|
|
|
class TestSearchsorted:
|
|
@pytest.mark.parametrize("freq", ["D", "2D"])
|
|
def test_searchsorted(self, freq):
|
|
pidx = PeriodIndex(
|
|
["2014-01-01", "2014-01-02", "2014-01-03", "2014-01-04", "2014-01-05"],
|
|
freq=freq,
|
|
)
|
|
|
|
p1 = Period("2014-01-01", freq=freq)
|
|
assert pidx.searchsorted(p1) == 0
|
|
|
|
p2 = Period("2014-01-04", freq=freq)
|
|
assert pidx.searchsorted(p2) == 3
|
|
|
|
if np_version_under1p18:
|
|
# GH#36254
|
|
# Following numpy convention, NaT goes at the beginning
|
|
# (unlike NaN which goes at the end)
|
|
assert pidx.searchsorted(NaT) == 0
|
|
else:
|
|
assert pidx.searchsorted(NaT) == 5
|
|
|
|
msg = "Input has different freq=H from PeriodArray"
|
|
with pytest.raises(IncompatibleFrequency, match=msg):
|
|
pidx.searchsorted(Period("2014-01-01", freq="H"))
|
|
|
|
msg = "Input has different freq=5D from PeriodArray"
|
|
with pytest.raises(IncompatibleFrequency, match=msg):
|
|
pidx.searchsorted(Period("2014-01-01", freq="5D"))
|
|
|
|
@pytest.mark.parametrize("klass", [list, np.array, array, Series])
|
|
def test_searchsorted_different_argument_classes(self, klass):
|
|
pidx = PeriodIndex(
|
|
["2014-01-01", "2014-01-02", "2014-01-03", "2014-01-04", "2014-01-05"],
|
|
freq="D",
|
|
)
|
|
result = pidx.searchsorted(klass(pidx))
|
|
expected = np.arange(len(pidx), dtype=result.dtype)
|
|
tm.assert_numpy_array_equal(result, expected)
|
|
|
|
result = pidx._data.searchsorted(klass(pidx))
|
|
tm.assert_numpy_array_equal(result, expected)
|
|
|
|
def test_searchsorted_invalid(self):
|
|
pidx = PeriodIndex(
|
|
["2014-01-01", "2014-01-02", "2014-01-03", "2014-01-04", "2014-01-05"],
|
|
freq="D",
|
|
)
|
|
|
|
other = np.array([0, 1], dtype=np.int64)
|
|
|
|
msg = "|".join(
|
|
[
|
|
"searchsorted requires compatible dtype or scalar",
|
|
"value should be a 'Period', 'NaT', or array of those. Got",
|
|
]
|
|
)
|
|
with pytest.raises(TypeError, match=msg):
|
|
pidx.searchsorted(other)
|
|
|
|
with pytest.raises(TypeError, match=msg):
|
|
pidx.searchsorted(other.astype("timedelta64[ns]"))
|
|
|
|
with pytest.raises(TypeError, match=msg):
|
|
pidx.searchsorted(np.timedelta64(4))
|
|
|
|
with pytest.raises(TypeError, match=msg):
|
|
pidx.searchsorted(np.timedelta64("NaT", "ms"))
|
|
|
|
with pytest.raises(TypeError, match=msg):
|
|
pidx.searchsorted(np.datetime64(4, "ns"))
|
|
|
|
with pytest.raises(TypeError, match=msg):
|
|
pidx.searchsorted(np.datetime64("NaT", "ns"))
|