39
39
ScalarIndexer ,
40
40
SequenceIndexer ,
41
41
SortKind ,
42
+ TimeArrayLike ,
42
43
npt ,
43
44
)
44
45
from pandas .compat .numpy import function as nv
78
79
unique ,
79
80
value_counts ,
80
81
)
82
+ from pandas .core .arrays .datetimes import DatetimeArray
83
+ from pandas .core .arrays .timedeltas import TimedeltaArray
81
84
from pandas .core .arrays .base import (
82
85
ExtensionArray ,
83
86
_extension_array_shared_docs ,
102
105
103
106
104
107
IntervalArrayT = TypeVar ("IntervalArrayT" , bound = "IntervalArray" )
108
+ IntervalSideT = Union [TimeArrayLike , np .ndarray ]
105
109
IntervalOrNA = Union [Interval , float ]
106
110
107
111
_interval_shared_docs : dict [str , str ] = {}
123
127
Parameters
124
128
----------
125
129
data : array-like (1-dimensional)
126
- Array-like containing Interval objects from which to build the
127
- %(klass)s.
130
+ Array-like (ndarray, :class:`DateTimeArray`, :class:`TimeDeltaArray`) containing
131
+ Interval objects from which to build the %(klass)s.
128
132
closed : {'left', 'right', 'both', 'neither'}, default 'right'
129
133
Whether the intervals are closed on the left-side, right-side, both or
130
134
neither.
@@ -213,8 +217,8 @@ def ndim(self) -> Literal[1]:
213
217
return 1
214
218
215
219
# To make mypy recognize the fields
216
- _left : np . ndarray
217
- _right : np . ndarray
220
+ _left : IntervalSideT
221
+ _right : IntervalSideT
218
222
_dtype : IntervalDtype
219
223
220
224
# ---------------------------------------------------------------------
@@ -232,8 +236,8 @@ def __new__(
232
236
data = extract_array (data , extract_numpy = True )
233
237
234
238
if isinstance (data , cls ):
235
- left : ArrayLike = data ._left
236
- right : ArrayLike = data ._right
239
+ left : IntervalSideT = data ._left
240
+ right : IntervalSideT = data ._right
237
241
closed = closed or data .closed
238
242
dtype = IntervalDtype (left .dtype , closed = closed )
239
243
else :
@@ -276,8 +280,8 @@ def __new__(
276
280
@classmethod
277
281
def _simple_new (
278
282
cls : type [IntervalArrayT ],
279
- left ,
280
- right ,
283
+ left : IntervalSideT ,
284
+ right : IntervalSideT ,
281
285
dtype : IntervalDtype ,
282
286
) -> IntervalArrayT :
283
287
result = IntervalMixin .__new__ (cls )
@@ -295,7 +299,7 @@ def _ensure_simple_new_inputs(
295
299
closed : IntervalClosedType | None = None ,
296
300
copy : bool = False ,
297
301
dtype : Dtype | None = None ,
298
- ) -> tuple [ArrayLike , ArrayLike , IntervalDtype ]:
302
+ ) -> tuple [IntervalSideT , IntervalSideT , IntervalDtype ]:
299
303
"""Ensure correctness of input parameters for cls._simple_new."""
300
304
from pandas .core .indexes .base import ensure_index
301
305
@@ -1574,9 +1578,11 @@ def _putmask(self, mask: npt.NDArray[np.bool_], value) -> None:
1574
1578
1575
1579
if isinstance (self ._left , np .ndarray ):
1576
1580
np .putmask (self ._left , mask , value_left )
1581
+ assert isinstance (self ._right , np .ndarray )
1577
1582
np .putmask (self ._right , mask , value_right )
1578
1583
else :
1579
1584
self ._left ._putmask (mask , value_left )
1585
+ assert not isinstance (self ._right , np .ndarray )
1580
1586
self ._right ._putmask (mask , value_right )
1581
1587
1582
1588
def insert (self : IntervalArrayT , loc : int , item : Interval ) -> IntervalArrayT :
@@ -1604,9 +1610,11 @@ def insert(self: IntervalArrayT, loc: int, item: Interval) -> IntervalArrayT:
1604
1610
def delete (self : IntervalArrayT , loc ) -> IntervalArrayT :
1605
1611
if isinstance (self ._left , np .ndarray ):
1606
1612
new_left = np .delete (self ._left , loc )
1613
+ assert isinstance (self ._right , np .ndarray )
1607
1614
new_right = np .delete (self ._right , loc )
1608
1615
else :
1609
1616
new_left = self ._left .delete (loc )
1617
+ assert not isinstance (self ._right , np .ndarray )
1610
1618
new_right = self ._right .delete (loc )
1611
1619
return self ._shallow_copy (left = new_left , right = new_right )
1612
1620
@@ -1707,7 +1715,7 @@ def isin(self, values) -> npt.NDArray[np.bool_]:
1707
1715
return isin (self .astype (object ), values .astype (object ))
1708
1716
1709
1717
@property
1710
- def _combined (self ) -> ArrayLike :
1718
+ def _combined (self ) -> IntervalSideT :
1711
1719
left = self .left ._values .reshape (- 1 , 1 )
1712
1720
right = self .right ._values .reshape (- 1 , 1 )
1713
1721
if needs_i8_conversion (left .dtype ):
@@ -1724,15 +1732,12 @@ def _from_combined(self, combined: np.ndarray) -> IntervalArray:
1724
1732
1725
1733
dtype = self ._left .dtype
1726
1734
if needs_i8_conversion (dtype ):
1727
- # error: "Type[ndarray[Any, Any]]" has no attribute "_from_sequence"
1728
- new_left = type (self ._left )._from_sequence ( # type: ignore[attr-defined]
1729
- nc [:, 0 ], dtype = dtype
1730
- )
1731
- # error: "Type[ndarray[Any, Any]]" has no attribute "_from_sequence"
1732
- new_right = type (self ._right )._from_sequence ( # type: ignore[attr-defined]
1733
- nc [:, 1 ], dtype = dtype
1734
- )
1735
+ assert isinstance (self ._left , (DatetimeArray , TimedeltaArray ))
1736
+ new_left = type (self ._left )._from_sequence (nc [:, 0 ], dtype = dtype )
1737
+ assert isinstance (self ._right , (DatetimeArray , TimedeltaArray ))
1738
+ new_right = type (self ._right )._from_sequence (nc [:, 1 ], dtype = dtype )
1735
1739
else :
1740
+ assert isinstance (dtype , np .dtype )
1736
1741
new_left = nc [:, 0 ].view (dtype )
1737
1742
new_right = nc [:, 1 ].view (dtype )
1738
1743
return self ._shallow_copy (left = new_left , right = new_right )
0 commit comments