Skip to content

Commit 5911bed

Browse files
cmp0xffcmp0xff
committed
refactor after separating dtype args
1 parent c12136f commit 5911bed

File tree

5 files changed

+86
-111
lines changed

5 files changed

+86
-111
lines changed

‎pandas-stubs/core/construction.pyi‎

Lines changed: 69 additions & 55 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ from pandas.core.arrays.numpy_ import NumpyExtensionArray
2121
from pandas.core.arrays.period import PeriodArray
2222
from pandas.core.arrays.sparse.array import SparseArray
2323
from pandas.core.arrays.string_ import StringArray
24+
from pandas.core.arrays.string_arrow import ArrowStringArray
2425
from pandas.core.arrays.timedeltas import TimedeltaArray
2526
from pandas.core.indexes.base import Index
2627
from pandas.core.indexes.category import CategoricalIndex
@@ -39,16 +40,22 @@ from pandas._libs.tslibs.period import Period
3940
from pandas._libs.tslibs.timedeltas import Timedelta
4041
from pandas._libs.tslibs.timestamps import Timestamp
4142
from pandas._typing import (
42-
BooleanDtypeArg,
43+
BuiltinBooleanDtypeArg,
44+
BuiltinFloatDtypeArg,
45+
BuiltinIntDtypeArg,
46+
BuiltinStrDtypeArg,
4347
CategoryDtypeArg,
44-
FloatDtypeArg,
45-
IntDtypeArg,
4648
IntervalT,
49+
NumpyDtypeArg,
50+
PandasBooleanDtypeArg,
51+
PandasFloatDtypeArg,
52+
PandasIntDtypeArg,
53+
PandasStrDtypeArg,
54+
PandasTimestampDtypeArg,
55+
PandasUIntDtypeArg,
56+
PyArrowNotStrDtypeArg,
57+
PyArrowStrDtypeArg,
4758
SequenceNotStr,
48-
StrDtypeArg,
49-
TimedeltaDtypeArg,
50-
TimestampDtypeArg,
51-
UIntDtypeArg,
5259
np_ndarray,
5360
np_ndarray_anyint,
5461
np_ndarray_bool,
@@ -64,18 +71,46 @@ from pandas.core.dtypes.dtypes import (
6471
PeriodDtype,
6572
)
6673

74+
# @overload
75+
# def array(
76+
# data: SequenceNotStr[NAType | None],
77+
# dtype: None = None,
78+
# copy: bool = True
79+
# ) -> NumpyExtensionArray: ...
6780
@overload
68-
def array( # type: ignore[overload-overlap]
81+
def array( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload]
6982
data: SequenceNotStr[Any] | np_ndarray | ExtensionArray | Index | Series,
7083
dtype: CategoryDtypeArg,
7184
copy: bool = True,
7285
) -> Categorical: ...
7386
@overload
87+
def array(
88+
# TODO: Categorical Series pandas-dev/pandas-stubs#1415
89+
data: Categorical | CategoricalIndex,
90+
dtype: CategoryDtypeArg | None = None,
91+
copy: bool = True,
92+
) -> Categorical: ...
93+
@overload
7494
def array( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload]
75-
data: Sequence[NAType | None],
95+
data: (
96+
Sequence[Period | NaTType | None] | PeriodArray | PeriodIndex | Series[Period]
97+
),
98+
dtype: PeriodDtype | None = None,
99+
copy: bool = True,
100+
) -> PeriodArray: ...
101+
@overload
102+
def array( # type: ignore[overload-overlap]
103+
# float("nan") also works, but I don't know how to put it in
104+
data: Sequence[IntervalT | None] | IntervalArray | IntervalIndex | Series[Interval],
105+
dtype: IntervalDtype | None = None,
106+
copy: bool = True,
107+
) -> IntervalArray: ...
108+
@overload
109+
def array(
110+
data: SparseArray | SparseIndex,
76111
dtype: str | np.dtype | ExtensionDtype | None = None,
77112
copy: bool = True,
78-
) -> NumpyExtensionArray: ...
113+
) -> SparseArray: ...
79114
@overload
80115
def array( # type: ignore[overload-overlap]
81116
data: (
@@ -85,21 +120,21 @@ def array( # type: ignore[overload-overlap]
85120
| Index[bool]
86121
| Series[int]
87122
),
88-
dtype: BooleanDtypeArg | None = None,
123+
dtype: BuiltinBooleanDtypeArg | PandasBooleanDtypeArg | None = None,
89124
copy: bool = True,
90125
) -> BooleanArray: ...
91126
@overload
92127
def array( # type: ignore[overload-overlap]
93128
data: Sequence[int | np.integer | NAType | None] | np_ndarray_anyint | IntegerArray,
94-
dtype: IntDtypeArg | UIntDtypeArg | None = None,
129+
dtype: BuiltinIntDtypeArg | PandasIntDtypeArg | PandasUIntDtypeArg | None = None,
95130
copy: bool = True,
96131
) -> IntegerArray: ...
97132
@overload
98133
def array( # type: ignore[overload-overlap]
99134
data: (
100135
Sequence[float | np.floating | NAType | None] | np_ndarray_float | FloatingArray
101136
),
102-
dtype: FloatDtypeArg | None = None,
137+
dtype: BuiltinFloatDtypeArg | PandasFloatDtypeArg | None = None,
103138
copy: bool = True,
104139
) -> FloatingArray: ...
105140
@overload
@@ -112,75 +147,54 @@ def array( # type: ignore[overload-overlap]
112147
| DatetimeIndex
113148
| Series[Timestamp]
114149
),
115-
dtype: TimestampDtypeArg | None = None,
150+
dtype: PandasTimestampDtypeArg | None = None,
116151
copy: bool = True,
117152
) -> DatetimeArray: ...
118153
@overload
119-
def array( # type: ignore[overload-overlap]
154+
def array(
120155
data: (
121156
Sequence[timedelta | np.timedelta64 | NaTType | None]
122157
| np_ndarray_td
123158
| TimedeltaArray
124159
| TimedeltaIndex
125160
| Series[Timedelta]
126161
),
127-
dtype: TimedeltaDtypeArg | None = None,
162+
dtype: None = None,
128163
copy: bool = True,
129164
) -> TimedeltaArray: ...
130165
@overload
131-
def array(
166+
def array( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload]
132167
data: SequenceNotStr[str | np.str_ | NAType | None] | np_ndarray_str | StringArray,
133-
dtype: StrDtypeArg | None = None,
168+
dtype: BuiltinStrDtypeArg | PandasStrDtypeArg | None = None,
134169
copy: bool = True,
135170
) -> StringArray: ...
136171
@overload
137172
def array( # type: ignore[overload-overlap]
138173
data: (
139-
Sequence[Period | NaTType | None] | PeriodArray | PeriodIndex | Series[Period]
174+
SequenceNotStr[str | np.str_ | NAType | None]
175+
| np_ndarray_str
176+
| StringArray
177+
| ArrowStringArray
140178
),
141-
dtype: PeriodDtype | None = None,
142-
copy: bool = True,
143-
) -> PeriodArray: ...
144-
@overload
145-
def array(
146-
# float("nan") also works, but I don't know how to put it in
147-
data: Sequence[IntervalT | None] | IntervalArray | IntervalIndex | Series[Interval],
148-
dtype: IntervalDtype | None = None,
179+
dtype: PyArrowStrDtypeArg | None = None,
149180
copy: bool = True,
150-
) -> IntervalArray: ...
151-
@overload
152-
def array(
153-
# TODO: Categorical Series pandas-dev/pandas-stubs#1415
154-
data: Categorical | CategoricalIndex,
155-
dtype: CategoryDtypeArg | None = None,
156-
copy: bool = True,
157-
) -> Categorical: ...
181+
) -> ArrowStringArray: ...
158182
@overload
159183
def array(
160-
data: (
161-
SequenceNotStr[object]
162-
| np.typing.NDArray[np.object_]
163-
| NumpyExtensionArray
164-
| RangeIndex
165-
),
166-
dtype: str | np.dtype | ExtensionDtype | None = None,
184+
data: SequenceNotStr[object] | np_ndarray | NumpyExtensionArray | RangeIndex,
185+
dtype: NumpyDtypeArg | None = None,
167186
copy: bool = True,
168187
) -> NumpyExtensionArray: ...
169188
@overload
170-
def array(
171-
data: SparseArray | SparseIndex,
172-
dtype: str | np.dtype | ExtensionDtype | None = None,
173-
copy: bool = True,
174-
) -> SparseArray: ...
175-
@overload
176189
def array(
177190
data: ArrowExtensionArray,
178-
dtype: str | np.dtype | ExtensionDtype | None = None,
191+
dtype: PyArrowNotStrDtypeArg | None = None,
179192
copy: bool = True,
180193
) -> ArrowExtensionArray: ...
181-
@overload
182-
def array(
183-
data: SequenceNotStr[Any] | np_ndarray | ExtensionArray | Index | Series,
184-
dtype: str | np.dtype | ExtensionDtype | None = None,
185-
copy: bool = True,
186-
) -> ExtensionArray: ...
194+
195+
# @overload
196+
# def array(
197+
# data: SequenceNotStr[Any] | np_ndarray | ExtensionArray | Index | Series,
198+
# dtype: str | np.dtype | ExtensionDtype | None = None,
199+
# copy: bool = True,
200+
# ) -> ExtensionArray: ...

‎tests/arrays/test_categorical.py‎

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -9,9 +9,7 @@
99
def test_constructor() -> None:
1010
check(assert_type(pd.array(["🐼"], dtype="category"), Categorical), Categorical)
1111
check(
12-
assert_type( # type: ignore[assert-type] # I do not understand
13-
pd.array(np.array(["🐼"]), dtype="category"), Categorical
14-
),
12+
assert_type(pd.array(np.array(["🐼"]), dtype="category"), Categorical),
1513
Categorical,
1614
)
1715
check(
@@ -32,9 +30,7 @@ def test_constructor() -> None:
3230
Categorical,
3331
)
3432
check(
35-
assert_type( # type: ignore[assert-type] # I do not understand
36-
pd.array(pd.Index(["🐼"], dtype="category")), Categorical
37-
),
33+
assert_type(pd.array(pd.Index(["🐼"], dtype="category")), Categorical),
3834
Categorical,
3935
)
4036
# TODO: Categorical Series pandas-dev/pandas-stubs#1415

‎tests/arrays/test_interval_array.py‎

Lines changed: 3 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -7,25 +7,10 @@
77

88
def test_constructor() -> None:
99
itv = pd.Interval(0, 1)
10-
check(
11-
assert_type( # type: ignore[assert-type] # I do not understand
12-
pd.array([itv]), IntervalArray
13-
),
14-
IntervalArray,
15-
)
16-
check(
17-
assert_type( # type: ignore[assert-type] # I do not understand
18-
pd.array([itv, None]), IntervalArray
19-
),
20-
IntervalArray,
21-
)
10+
check(assert_type(pd.array([itv]), IntervalArray), IntervalArray)
11+
check(assert_type(pd.array([itv, None]), IntervalArray), IntervalArray)
2212

23-
check(
24-
assert_type( # type: ignore[assert-type] # I do not understand
25-
pd.array(pd.array([itv])), IntervalArray
26-
),
27-
IntervalArray,
28-
)
13+
check(assert_type(pd.array(pd.array([itv])), IntervalArray), IntervalArray)
2914

3015
check(assert_type(pd.array(pd.Index([itv])), IntervalArray), IntervalArray)
3116

‎tests/arrays/test_numpy_extension_array.py‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -7,9 +7,9 @@
77

88

99
def test_constructor() -> None:
10-
check(
11-
assert_type(pd.array([pd.NA, None]), NumpyExtensionArray), NumpyExtensionArray
12-
)
10+
# check(
11+
# assert_type(pd.array([pd.NA, None]), NumpyExtensionArray), NumpyExtensionArray
12+
# )
1313

1414
check(
1515
assert_type( # type: ignore[assert-type] # I do not understand
@@ -18,15 +18,15 @@ def test_constructor() -> None:
1818
NumpyExtensionArray,
1919
)
2020
check(
21-
assert_type( # type: ignore[assert-type] # I do not understand, mypy must have problem with two Generic Variables somehow
21+
assert_type( # type: ignore[assert-type] # I do not understand
2222
pd.array(np.array([1, "🐼"], np.object_)), NumpyExtensionArray
2323
),
2424
NumpyExtensionArray,
2525
)
26-
check(
27-
assert_type(pd.array(pd.array([pd.NA, None])), NumpyExtensionArray),
28-
NumpyExtensionArray,
29-
)
26+
# check(
27+
# assert_type(pd.array(pd.array([pd.NA, None])), NumpyExtensionArray),
28+
# NumpyExtensionArray,
29+
# )
3030
check(
3131
assert_type(pd.array(pd.RangeIndex(0, 1)), NumpyExtensionArray),
3232
NumpyExtensionArray,

‎tests/arrays/test_period_array.py‎

Lines changed: 4 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -7,31 +7,11 @@
77

88
def test_constructor() -> None:
99
prd = pd.Period("2023-01-01")
10-
check(
11-
assert_type( # type: ignore[assert-type] # I do not understand
12-
pd.array([prd]), PeriodArray
13-
),
14-
PeriodArray,
15-
)
16-
check(
17-
assert_type( # type: ignore[assert-type] # I do not understand
18-
pd.array([prd, None]), PeriodArray
19-
),
20-
PeriodArray,
21-
)
22-
check(
23-
assert_type( # type: ignore[assert-type] # I do not understand
24-
pd.array([prd, pd.NaT, None]), PeriodArray
25-
),
26-
PeriodArray,
27-
)
10+
check(assert_type(pd.array([prd]), PeriodArray), PeriodArray)
11+
check(assert_type(pd.array([prd, None]), PeriodArray), PeriodArray)
12+
check(assert_type(pd.array([prd, pd.NaT, None]), PeriodArray), PeriodArray)
2813

29-
check(
30-
assert_type( # type: ignore[assert-type] # I do not understand
31-
pd.array(pd.array([prd])), PeriodArray
32-
),
33-
PeriodArray,
34-
)
14+
check(assert_type(pd.array(pd.array([prd])), PeriodArray), PeriodArray)
3515

3616
check(assert_type(pd.array(pd.Index([prd])), PeriodArray), PeriodArray)
3717

0 commit comments

Comments
 (0)