Skip to content

Commit 83675db

Browse files
cmp0xffcmp0xff
committed
feat: add ExtensionArrays
1 parent b54625a commit 83675db

File tree

12 files changed

+378
-45
lines changed

12 files changed

+378
-45
lines changed

‎pandas-stubs/core/algorithms.pyi‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -39,10 +39,10 @@ def unique(values: IntervalIndex[IntervalT]) -> IntervalIndex[IntervalT]: ...
3939
@overload
4040
def unique(values: PeriodIndex) -> PeriodIndex: ...
4141
@overload
42-
# switch to DatetimeIndex after Pandas 3.0
42+
# TODO: switch to DatetimeIndex after Pandas 3.0 pandas-dev/pandas#57064
4343
def unique(values: DatetimeIndex) -> np_1darray_dt | DatetimeIndex: ...
4444
@overload
45-
# switch to TimedeltaIndex after Pandas 3.0
45+
# TODO: switch to TimedeltaIndex after Pandas 3.0 pandas-dev/pandas#57064
4646
def unique(values: TimedeltaIndex) -> np_1darray_td: ...
4747
@overload
4848
# switch to Index[int] after Pandas 3.0
@@ -59,7 +59,8 @@ else:
5959
) -> np_ndarray: ...
6060

6161
@overload
62-
def unique(values: Index) -> np_1darray | Index: ... # switch to Index after Pandas 3.0
62+
# TODO: switch to Index after Pandas 3.0 pandas-dev/pandas#57064
63+
def unique(values: Index) -> np_1darray | Index: ...
6364
@overload
6465
def unique(values: Categorical) -> Categorical: ...
6566

‎pandas-stubs/core/arrays/__init__.pyi‎

Lines changed: 33 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,33 @@
1-
from pandas.core.arrays.base import ExtensionArray as ExtensionArray
2-
from pandas.core.arrays.boolean import BooleanArray as BooleanArray
3-
from pandas.core.arrays.categorical import Categorical as Categorical
4-
from pandas.core.arrays.datetimes import DatetimeArray as DatetimeArray
5-
from pandas.core.arrays.integer import IntegerArray as IntegerArray
6-
from pandas.core.arrays.interval import IntervalArray as IntervalArray
7-
from pandas.core.arrays.numpy_ import NumpyExtensionArray as NumpyExtensionArray
8-
from pandas.core.arrays.period import PeriodArray as PeriodArray
9-
from pandas.core.arrays.sparse import SparseArray as SparseArray
10-
from pandas.core.arrays.string_ import StringArray as StringArray
11-
from pandas.core.arrays.timedeltas import TimedeltaArray as TimedeltaArray
1+
from pandas.core.arrays.arrow import ArrowExtensionArray
2+
from pandas.core.arrays.base import ExtensionArray
3+
from pandas.core.arrays.boolean import BooleanArray
4+
from pandas.core.arrays.categorical import Categorical
5+
from pandas.core.arrays.datetimes import DatetimeArray
6+
from pandas.core.arrays.floating import FloatingArray
7+
from pandas.core.arrays.integer import IntegerArray
8+
from pandas.core.arrays.interval import IntervalArray
9+
from pandas.core.arrays.masked import BaseMaskedArray
10+
from pandas.core.arrays.numpy_ import NumpyExtensionArray
11+
from pandas.core.arrays.period import PeriodArray
12+
from pandas.core.arrays.sparse import SparseArray
13+
from pandas.core.arrays.string_ import StringArray
14+
from pandas.core.arrays.string_arrow import ArrowStringArray
15+
from pandas.core.arrays.timedeltas import TimedeltaArray
16+
17+
__all__ = [
18+
"ArrowExtensionArray",
19+
"ArrowStringArray",
20+
"BaseMaskedArray",
21+
"BooleanArray",
22+
"Categorical",
23+
"DatetimeArray",
24+
"ExtensionArray",
25+
"FloatingArray",
26+
"IntegerArray",
27+
"IntervalArray",
28+
"NumpyExtensionArray",
29+
"PeriodArray",
30+
"SparseArray",
31+
"StringArray",
32+
"TimedeltaArray",
33+
]

‎pandas-stubs/core/arrays/datetimes.pyi‎

Lines changed: 19 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,5 @@
11
from collections.abc import Iterator
22
from datetime import tzinfo as _tzinfo
3-
import sys
4-
from typing import Any
53

64
import numpy as np
75
from pandas.core.arrays.datetimelike import (
@@ -10,44 +8,44 @@ from pandas.core.arrays.datetimelike import (
108
TimelikeOps,
119
)
1210
from pandas.core.arrays.period import PeriodArray
13-
from pandas.core.indexes.base import Index
11+
from pandas.core.indexes.datetimes import DatetimeIndex
1412
from pandas.core.series import Series
15-
from typing_extensions import (
16-
Self,
17-
)
13+
from typing_extensions import Self
1814

1915
from pandas._libs.tslibs.timestamps import Timestamp
2016
from pandas._typing import (
21-
DtypeArg,
2217
Frequency,
18+
NumpyTimestampDtypeArg,
19+
PandasTimestampDtypeArg,
2320
TimeAmbiguous,
2421
TimeNonexistent,
2522
TimeZones,
2623
np_1darray,
2724
np_1darray_float,
2825
np_1darray_object,
2926
np_ndarray,
27+
np_ndarray_dt,
3028
)
3129

3230
from pandas.core.dtypes.dtypes import DatetimeTZDtype as DatetimeTZDtype
3331

3432
class DatetimeArray(DatetimeLikeArrayMixin, TimelikeOps, DatelikeOps):
3533
__array_priority__: int = ...
36-
def __init__(
37-
self,
38-
values: Series | Index | DatetimeArray | np_ndarray,
39-
dtype: DtypeArg | None = None,
40-
freq: Frequency | None = None,
34+
def __new__(
35+
cls,
36+
values: np_ndarray_dt | DatetimeIndex | Series[Timestamp] | Self,
37+
dtype: (
38+
np.dtypes.DateTime64DType
39+
| NumpyTimestampDtypeArg
40+
| PandasTimestampDtypeArg
41+
| None
42+
) = None,
4143
copy: bool = False,
42-
) -> None: ...
43-
# ignore in dtype() is from the pandas source
44-
if sys.version_info >= (3, 11):
45-
@property
46-
def dtype(self) -> np.dtype | DatetimeTZDtype: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
47-
else:
48-
@property
49-
def dtype(self) -> np.dtype[Any] | DatetimeTZDtype: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride] # pyrefly: ignore[bad-override]
50-
44+
) -> Self: ...
45+
@property
46+
def dtype( # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride] # pyrefly: ignore[bad-override]
47+
self,
48+
) -> np.dtypes.DateTime64DType | DatetimeTZDtype: ...
5149
@property
5250
def tz(self) -> _tzinfo | None: ...
5351
@property

‎pandas-stubs/core/arrays/integer.pyi‎

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
from pandas.core.arrays.masked import BaseMaskedArray
2+
from typing_extensions import Self
23

34
from pandas._libs.missing import NAType
45
from pandas._typing import (
@@ -20,12 +21,9 @@ class _IntegerDtype(ExtensionDtype):
2021
class IntegerArray(BaseMaskedArray):
2122
@property
2223
def dtype(self) -> _IntegerDtype: ...
23-
def __init__(
24-
self,
25-
values: AnyArrayLike,
26-
mask: np_ndarray_bool,
27-
copy: bool = False,
28-
) -> None: ...
24+
def __new__(
25+
cls, values: AnyArrayLike, mask: np_ndarray_bool, copy: bool = False
26+
) -> Self: ...
2927

3028
class Int8Dtype(_IntegerDtype): ...
3129
class Int16Dtype(_IntegerDtype): ...

‎pandas-stubs/core/construction.pyi‎

Lines changed: 108 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,43 +1,80 @@
11
from collections.abc import Sequence
2+
from datetime import (
3+
datetime,
4+
timedelta,
5+
)
6+
import sys
27
from typing import (
38
Any,
49
TypeAlias,
510
overload,
611
)
712

813
import numpy as np
14+
from pandas.core.arrays.base import ExtensionArray
915
from pandas.core.arrays.boolean import BooleanArray
16+
from pandas.core.arrays.categorical import Categorical
17+
from pandas.core.arrays.datetimes import DatetimeArray
1018
from pandas.core.arrays.floating import FloatingArray
1119
from pandas.core.arrays.integer import IntegerArray
20+
from pandas.core.arrays.interval import IntervalArray
1221
from pandas.core.arrays.numpy_ import NumpyExtensionArray
22+
from pandas.core.arrays.period import PeriodArray
23+
from pandas.core.arrays.sparse.array import SparseArray
1324
from pandas.core.arrays.string_ import (
1425
BaseStringArray,
1526
StringArray,
1627
StringDtype,
1728
)
1829
from pandas.core.arrays.string_arrow import ArrowStringArray
30+
from pandas.core.arrays.timedeltas import TimedeltaArray
31+
from pandas.core.indexes.base import Index
32+
from pandas.core.indexes.category import CategoricalIndex
33+
from pandas.core.indexes.datetimes import DatetimeIndex
34+
from pandas.core.indexes.interval import IntervalIndex
35+
from pandas.core.indexes.period import PeriodIndex
1936
from pandas.core.indexes.range import RangeIndex
37+
from pandas.core.indexes.timedeltas import TimedeltaIndex
38+
from pandas.core.series import Series
2039
from typing_extensions import Never
2140

41+
from pandas._libs.interval import Interval
2242
from pandas._libs.missing import NAType
43+
from pandas._libs.sparse import SparseIndex
2344
from pandas._libs.tslibs.nattype import NaTType
45+
from pandas._libs.tslibs.period import Period
46+
from pandas._libs.tslibs.timedeltas import Timedelta
47+
from pandas._libs.tslibs.timestamps import Timestamp
2448
from pandas._typing import (
2549
BuiltinDtypeArg,
50+
CategoryDtypeArg,
51+
IntervalT,
2652
Just,
2753
NumpyNotTimeDtypeArg,
54+
NumpyTimedeltaDtypeArg,
55+
NumpyTimestampDtypeArg,
2856
PandasBaseStrDtypeArg,
2957
PandasBooleanDtypeArg,
3058
PandasFloatDtypeArg,
3159
PandasIntDtypeArg,
3260
PandasStrDtypeArg,
61+
PandasTimestampDtypeArg,
3362
PandasUIntDtypeArg,
3463
PyArrowStrDtypeArg,
3564
SequenceNotStr,
3665
np_ndarray,
3766
np_ndarray_anyint,
3867
np_ndarray_bool,
68+
np_ndarray_dt,
3969
np_ndarray_float,
4070
np_ndarray_str,
71+
np_ndarray_td,
72+
)
73+
74+
from pandas.core.dtypes.base import ExtensionDtype
75+
from pandas.core.dtypes.dtypes import (
76+
IntervalDtype,
77+
PeriodDtype,
4178
)
4279

4380
_NaNNullableStrData: TypeAlias = (
@@ -62,6 +99,51 @@ def array( # type: ignore[overload-overlap] # pyright: ignore[reportOverlapping
6299
dtype: BuiltinDtypeArg | NumpyNotTimeDtypeArg | None = None,
63100
copy: bool = True,
64101
) -> NumpyExtensionArray: ...
102+
@overload
103+
def array( # type: ignore[overload-overlap]
104+
data: SequenceNotStr[Any] | np_ndarray | ExtensionArray | Index | Series,
105+
dtype: CategoryDtypeArg,
106+
copy: bool = True,
107+
) -> Categorical: ...
108+
@overload
109+
def array(
110+
# TODO: Categorical Series pandas-dev/pandas-stubs#1415
111+
data: Categorical | CategoricalIndex,
112+
dtype: CategoryDtypeArg | None = None,
113+
copy: bool = True,
114+
) -> Categorical: ...
115+
@overload
116+
def array( # type: ignore[overload-overlap]
117+
data: (
118+
Sequence[Period | NaTType | None] | PeriodArray | PeriodIndex | Series[Period]
119+
),
120+
dtype: PeriodDtype | None = None,
121+
copy: bool = True,
122+
) -> PeriodArray: ...
123+
@overload
124+
def array( # type: ignore[overload-overlap]
125+
# float("nan") also works, but I don't know how to put it in
126+
data: Sequence[IntervalT | None] | IntervalArray | IntervalIndex | Series[Interval],
127+
dtype: IntervalDtype | None = None,
128+
copy: bool = True,
129+
) -> IntervalArray: ...
130+
131+
if sys.version_info >= (3, 11):
132+
@overload
133+
def array(
134+
data: SparseArray | SparseIndex,
135+
dtype: str | np.dtype | ExtensionDtype | None = None,
136+
copy: bool = True,
137+
) -> SparseArray: ...
138+
139+
else:
140+
@overload
141+
def array(
142+
data: SparseArray | SparseIndex,
143+
dtype: str | np.dtype[Any] | ExtensionDtype | None = None,
144+
copy: bool = True,
145+
) -> SparseArray: ...
146+
65147
@overload
66148
def array(
67149
data: Sequence[bool | np.bool | Just[float] | NAType | None],
@@ -99,14 +181,39 @@ def array(
99181
copy: bool = True,
100182
) -> IntegerArray: ...
101183
@overload
102-
def array(
184+
def array( # type: ignore[overload-overlap]
103185
data: (
104186
Sequence[float | np.floating | NAType | None] | np_ndarray_float | FloatingArray
105187
),
106188
dtype: PandasFloatDtypeArg | None = None,
107189
copy: bool = True,
108190
) -> FloatingArray: ...
109191
@overload
192+
def array( # type: ignore[overload-overlap]
193+
data: ( # TODO: merge the two Sequence's after 3.0 pandas-dev/pandas#57064
194+
Sequence[datetime | NaTType | None]
195+
| Sequence[np.datetime64 | NaTType | None]
196+
| np_ndarray_dt
197+
| DatetimeArray
198+
| DatetimeIndex
199+
| Series[Timestamp]
200+
),
201+
dtype: PandasTimestampDtypeArg | NumpyTimestampDtypeArg | None = None,
202+
copy: bool = True,
203+
) -> DatetimeArray: ...
204+
@overload
205+
def array(
206+
data: (
207+
Sequence[timedelta | np.timedelta64 | NaTType | None]
208+
| np_ndarray_td
209+
| TimedeltaArray
210+
| TimedeltaIndex
211+
| Series[Timedelta]
212+
),
213+
dtype: NumpyTimedeltaDtypeArg | None = None,
214+
copy: bool = True,
215+
) -> TimedeltaArray: ...
216+
@overload
110217
def array( # type: ignore[overload-overlap]
111218
data: _NaNNullableStrData, dtype: StringDtype[Never], copy: bool = True
112219
) -> BaseStringArray: ...

‎tests/_typing.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
# ruff: noqa: PLC0414
12
# This file serves as a stub file for static type checkers
23
# (pyright does not like it if I call the file tests/_typing.pyi).
34
# By tests/conftest.py, pytest copies this file to tests/_typing.pyi, and

0 commit comments

Comments
 (0)