| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231 |
- import numpy as np
- import pytest
- import pandas as pd
- from pandas import (
- Index,
- Interval,
- IntervalIndex,
- Timedelta,
- Timestamp,
- date_range,
- timedelta_range,
- )
- import pandas._testing as tm
- from pandas.core.arrays import IntervalArray
- @pytest.fixture(
- params=[
- (Index([0, 2, 4]), Index([1, 3, 5])),
- (Index([0.0, 1.0, 2.0]), Index([1.0, 2.0, 3.0])),
- (timedelta_range("0 days", periods=3), timedelta_range("1 day", periods=3)),
- (date_range("20170101", periods=3), date_range("20170102", periods=3)),
- (
- date_range("20170101", periods=3, tz="US/Eastern"),
- date_range("20170102", periods=3, tz="US/Eastern"),
- ),
- ],
- ids=lambda x: str(x[0].dtype),
- )
- def left_right_dtypes(request):
- """
- Fixture for building an IntervalArray from various dtypes
- """
- return request.param
- class TestAttributes:
- @pytest.mark.parametrize(
- "left, right",
- [
- (0, 1),
- (Timedelta("0 days"), Timedelta("1 day")),
- (Timestamp("2018-01-01"), Timestamp("2018-01-02")),
- (
- Timestamp("2018-01-01", tz="US/Eastern"),
- Timestamp("2018-01-02", tz="US/Eastern"),
- ),
- ],
- )
- @pytest.mark.parametrize("constructor", [IntervalArray, IntervalIndex])
- def test_is_empty(self, constructor, left, right, closed):
- # GH27219
- tuples = [(left, left), (left, right), np.nan]
- expected = np.array([closed != "both", False, False])
- result = constructor.from_tuples(tuples, closed=closed).is_empty
- tm.assert_numpy_array_equal(result, expected)
- class TestMethods:
- @pytest.mark.parametrize("new_closed", ["left", "right", "both", "neither"])
- def test_set_closed(self, closed, new_closed):
- # GH 21670
- array = IntervalArray.from_breaks(range(10), closed=closed)
- result = array.set_closed(new_closed)
- expected = IntervalArray.from_breaks(range(10), closed=new_closed)
- tm.assert_extension_array_equal(result, expected)
- @pytest.mark.parametrize(
- "other",
- [
- Interval(0, 1, closed="right"),
- IntervalArray.from_breaks([1, 2, 3, 4], closed="right"),
- ],
- )
- def test_where_raises(self, other):
- # GH#45768 The IntervalArray methods raises; the Series method coerces
- ser = pd.Series(IntervalArray.from_breaks([1, 2, 3, 4], closed="left"))
- mask = np.array([True, False, True])
- match = "'value.closed' is 'right', expected 'left'."
- with pytest.raises(ValueError, match=match):
- ser.array._where(mask, other)
- res = ser.where(mask, other=other)
- expected = ser.astype(object).where(mask, other)
- tm.assert_series_equal(res, expected)
- def test_shift(self):
- # https://github.com/pandas-dev/pandas/issues/31495, GH#22428, GH#31502
- a = IntervalArray.from_breaks([1, 2, 3])
- result = a.shift()
- # int -> float
- expected = IntervalArray.from_tuples([(np.nan, np.nan), (1.0, 2.0)])
- tm.assert_interval_array_equal(result, expected)
- msg = "can only insert Interval objects and NA into an IntervalArray"
- with pytest.raises(TypeError, match=msg):
- a.shift(1, fill_value=pd.NaT)
- def test_shift_datetime(self):
- # GH#31502, GH#31504
- a = IntervalArray.from_breaks(date_range("2000", periods=4))
- result = a.shift(2)
- expected = a.take([-1, -1, 0], allow_fill=True)
- tm.assert_interval_array_equal(result, expected)
- result = a.shift(-1)
- expected = a.take([1, 2, -1], allow_fill=True)
- tm.assert_interval_array_equal(result, expected)
- msg = "can only insert Interval objects and NA into an IntervalArray"
- with pytest.raises(TypeError, match=msg):
- a.shift(1, fill_value=np.timedelta64("NaT", "ns"))
- class TestSetitem:
- def test_set_na(self, left_right_dtypes):
- left, right = left_right_dtypes
- left = left.copy(deep=True)
- right = right.copy(deep=True)
- result = IntervalArray.from_arrays(left, right)
- if result.dtype.subtype.kind not in ["m", "M"]:
- msg = "'value' should be an interval type, got <.*NaTType'> instead."
- with pytest.raises(TypeError, match=msg):
- result[0] = pd.NaT
- if result.dtype.subtype.kind in ["i", "u"]:
- msg = "Cannot set float NaN to integer-backed IntervalArray"
- # GH#45484 TypeError, not ValueError, matches what we get with
- # non-NA un-holdable value.
- with pytest.raises(TypeError, match=msg):
- result[0] = np.nan
- return
- result[0] = np.nan
- expected_left = Index([left._na_value] + list(left[1:]))
- expected_right = Index([right._na_value] + list(right[1:]))
- expected = IntervalArray.from_arrays(expected_left, expected_right)
- tm.assert_extension_array_equal(result, expected)
- def test_setitem_mismatched_closed(self):
- arr = IntervalArray.from_breaks(range(4))
- orig = arr.copy()
- other = arr.set_closed("both")
- msg = "'value.closed' is 'both', expected 'right'"
- with pytest.raises(ValueError, match=msg):
- arr[0] = other[0]
- with pytest.raises(ValueError, match=msg):
- arr[:1] = other[:1]
- with pytest.raises(ValueError, match=msg):
- arr[:0] = other[:0]
- with pytest.raises(ValueError, match=msg):
- arr[:] = other[::-1]
- with pytest.raises(ValueError, match=msg):
- arr[:] = list(other[::-1])
- with pytest.raises(ValueError, match=msg):
- arr[:] = other[::-1].astype(object)
- with pytest.raises(ValueError, match=msg):
- arr[:] = other[::-1].astype("category")
- # empty list should be no-op
- arr[:0] = []
- tm.assert_interval_array_equal(arr, orig)
- class TestReductions:
- def test_min_max_invalid_axis(self, left_right_dtypes):
- left, right = left_right_dtypes
- left = left.copy(deep=True)
- right = right.copy(deep=True)
- arr = IntervalArray.from_arrays(left, right)
- msg = "`axis` must be fewer than the number of dimensions"
- for axis in [-2, 1]:
- with pytest.raises(ValueError, match=msg):
- arr.min(axis=axis)
- with pytest.raises(ValueError, match=msg):
- arr.max(axis=axis)
- msg = "'>=' not supported between"
- with pytest.raises(TypeError, match=msg):
- arr.min(axis="foo")
- with pytest.raises(TypeError, match=msg):
- arr.max(axis="foo")
- def test_min_max(self, left_right_dtypes, index_or_series_or_array):
- # GH#44746
- left, right = left_right_dtypes
- left = left.copy(deep=True)
- right = right.copy(deep=True)
- arr = IntervalArray.from_arrays(left, right)
- # The expected results below are only valid if monotonic
- assert left.is_monotonic_increasing
- assert Index(arr).is_monotonic_increasing
- MIN = arr[0]
- MAX = arr[-1]
- indexer = np.arange(len(arr))
- np.random.default_rng(2).shuffle(indexer)
- arr = arr.take(indexer)
- arr_na = arr.insert(2, np.nan)
- arr = index_or_series_or_array(arr)
- arr_na = index_or_series_or_array(arr_na)
- for skipna in [True, False]:
- res = arr.min(skipna=skipna)
- assert res == MIN
- assert type(res) == type(MIN)
- res = arr.max(skipna=skipna)
- assert res == MAX
- assert type(res) == type(MAX)
- res = arr_na.min(skipna=False)
- assert np.isnan(res)
- res = arr_na.max(skipna=False)
- assert np.isnan(res)
- res = arr_na.min(skipna=True)
- assert res == MIN
- assert type(res) == type(MIN)
- res = arr_na.max(skipna=True)
- assert res == MAX
- assert type(res) == type(MAX)
|