Skip to content

Commit b5b9955

Browse files
jbrockmendelJulianWgs
jbrockmendel
authored and
JulianWgs
committed
TYP: libalgos.pyi (pandas-dev#40760)
1 parent 33222bd commit b5b9955

File tree

14 files changed

+331
-35
lines changed

14 files changed

+331
-35
lines changed

‎pandas/_libs/algos.pyi

+293
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,293 @@
1+
# Note: this covers algos.pyx and algos_common_helper but NOT algos_take_helper
2+
from typing import Any
3+
4+
import numpy as np
5+
6+
class Infinity:
7+
"""
8+
Provide a positive Infinity comparison method for ranking.
9+
"""
10+
def __eq__(self, other) -> bool: ...
11+
def __ne__(self, other) -> bool: ...
12+
def __lt__(self, other) -> bool: ...
13+
def __le__(self, other) -> bool: ...
14+
def __gt__(self, other) -> bool: ...
15+
def __ge__(self, other) -> bool: ...
16+
17+
18+
class NegInfinity:
19+
"""
20+
Provide a negative Infinity comparison method for ranking.
21+
"""
22+
def __eq__(self, other) -> bool: ...
23+
def __ne__(self, other) -> bool: ...
24+
def __lt__(self, other) -> bool: ...
25+
def __le__(self, other) -> bool: ...
26+
def __gt__(self, other) -> bool: ...
27+
def __ge__(self, other) -> bool: ...
28+
29+
30+
def unique_deltas(
31+
arr: np.ndarray, # const int64_t[:]
32+
) -> np.ndarray: ... # np.ndarray[np.int64, ndim=1]
33+
34+
35+
def is_lexsorted(list_of_arrays: list[np.ndarray]) -> bool: ...
36+
37+
38+
def groupsort_indexer(
39+
index: np.ndarray, # const int64_t[:]
40+
ngroups: int,
41+
) -> tuple[
42+
np.ndarray, # ndarray[int64_t, ndim=1]
43+
np.ndarray, # ndarray[int64_t, ndim=1]
44+
]:
45+
...
46+
47+
48+
def kth_smallest(
49+
a: np.ndarray, # numeric[:]
50+
k: int,
51+
) -> Any: ... # numeric
52+
53+
54+
# ----------------------------------------------------------------------
55+
# Pairwise correlation/covariance
56+
57+
58+
59+
def nancorr(
60+
mat: np.ndarray, # const float64_t[:, :]
61+
cov: bool = False,
62+
minp=None,
63+
) -> np.ndarray: # np.ndarray[float64_t, ndim=2]
64+
...
65+
66+
67+
def nancorr_spearman(
68+
mat: np.ndarray, # ndarray[float64_t, ndim=2]
69+
minp: int = 1,
70+
) -> np.ndarray: # np.ndarray[np.float64, ndim=2]
71+
...
72+
73+
74+
def nancorr_kendall(
75+
mat: np.ndarray, # ndarray[float64_t, ndim=2]
76+
minp: int = 1,
77+
) -> np.ndarray: # np.ndarray[float64, ndim=2]
78+
...
79+
80+
# ----------------------------------------------------------------------
81+
82+
# ctypedef fused algos_t:
83+
# float64_t
84+
# float32_t
85+
# object
86+
# int64_t
87+
# int32_t
88+
# int16_t
89+
# int8_t
90+
# uint64_t
91+
# uint32_t
92+
# uint16_t
93+
# uint8_t
94+
95+
96+
def validate_limit(nobs: int | None, limit=None) -> int: ...
97+
98+
99+
def pad(
100+
old: np.ndarray, # ndarray[algos_t]
101+
new: np.ndarray, # ndarray[algos_t]
102+
limit=None,
103+
) -> np.ndarray: ... # np.ndarray[np.intp, ndim=1]
104+
105+
106+
def pad_inplace(
107+
values: np.ndarray, # algos_t[:]
108+
mask: np.ndarray, # uint8_t[:]
109+
limit=None,
110+
) -> None: ...
111+
112+
113+
def pad_2d_inplace(
114+
values: np.ndarray, # algos_t[:, :]
115+
mask: np.ndarray, # const uint8_t[:, :]
116+
limit=None,
117+
) -> None:
118+
...
119+
120+
121+
def backfill(
122+
old: np.ndarray, # ndarray[algos_t]
123+
new: np.ndarray, # ndarray[algos_t]
124+
limit=None,
125+
) -> np.ndarray: # np.ndarray[np.intp, ndim=1]
126+
...
127+
128+
def backfill_inplace(
129+
values: np.ndarray, # algos_t[:]
130+
mask: np.ndarray, # uint8_t[:]
131+
limit=None,
132+
) -> None: ...
133+
134+
135+
def backfill_2d_inplace(
136+
values: np.ndarray, # algos_t[:, :]
137+
mask: np.ndarray, # const uint8_t[:, :]
138+
limit=None,
139+
) -> None: ...
140+
141+
142+
def is_monotonic(
143+
arr: np.ndarray, # ndarray[algos_t, ndim=1]
144+
timelike: bool
145+
) -> tuple[bool, bool, bool]:
146+
...
147+
148+
# ----------------------------------------------------------------------
149+
# rank_1d, rank_2d
150+
# ----------------------------------------------------------------------
151+
152+
# ctypedef fused rank_t:
153+
# object
154+
# float64_t
155+
# uint64_t
156+
# int64_t
157+
158+
159+
def rank_1d(
160+
values: np.ndarray, # ndarray[rank_t, ndim=1]
161+
labels: np.ndarray, # const int64_t[:]
162+
is_datetimelike: bool = ...,
163+
ties_method=...,
164+
ascending: bool = ...,
165+
pct: bool = ...,
166+
na_option=...,
167+
) -> np.ndarray: ... # np.ndarray[float64_t, ndim=1]
168+
169+
170+
def rank_2d(
171+
in_arr: np.ndarray, # ndarray[rank_t, ndim=2]
172+
axis: int = ...,
173+
is_datetimelike: bool = ...,
174+
ties_method=...,
175+
ascending: bool = ...,
176+
na_option=...,
177+
pct: bool = ...,
178+
) -> np.ndarray: ... # np.ndarray[float64_t, ndim=1]
179+
180+
181+
def diff_2d(
182+
arr: np.ndarray, # ndarray[diff_t, ndim=2]
183+
out: np.ndarray, # ndarray[out_t, ndim=2]
184+
periods: int,
185+
axis: int,
186+
datetimelike: bool = ...,
187+
) -> None: ...
188+
189+
190+
def ensure_platform_int(arr: object) -> np.ndarray: ...
191+
192+
def ensure_object(arr: object) -> np.ndarray: ...
193+
194+
def ensure_float64(arr: object, copy=True) -> np.ndarray: ...
195+
196+
def ensure_float32(arr: object, copy=True) -> np.ndarray: ...
197+
198+
def ensure_int8(arr: object, copy=True) -> np.ndarray: ...
199+
200+
def ensure_int16(arr: object, copy=True) -> np.ndarray: ...
201+
202+
def ensure_int32(arr: object, copy=True) -> np.ndarray: ...
203+
204+
def ensure_int64(arr: object, copy=True) -> np.ndarray: ...
205+
206+
def ensure_uint8(arr: object, copy=True) -> np.ndarray: ...
207+
208+
def ensure_uint16(arr: object, copy=True) -> np.ndarray: ...
209+
210+
def ensure_uint32(arr: object, copy=True) -> np.ndarray: ...
211+
212+
def ensure_uint64(arr: object, copy=True) -> np.ndarray: ...
213+
214+
215+
def take_1d_int8_int8(values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=...) -> None: ...
216+
def take_1d_int8_int32(values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=...) -> None: ...
217+
def take_1d_int8_int64(values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=...) -> None: ...
218+
def take_1d_int8_float64(values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=...) -> None: ...
219+
def take_1d_int16_int16(values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=...) -> None: ...
220+
def take_1d_int16_int32(values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=...) -> None: ...
221+
def take_1d_int16_int64(values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=...) -> None: ...
222+
def take_1d_int16_float64(values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=...) -> None: ...
223+
def take_1d_int32_int32(values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=...) -> None: ...
224+
def take_1d_int32_int64(values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=...) -> None: ...
225+
def take_1d_int32_float64(values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=...) -> None: ...
226+
def take_1d_int64_int64(values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=...) -> None: ...
227+
def take_1d_int64_float64(values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=...) -> None: ...
228+
def take_1d_float32_float32(values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=...) -> None: ...
229+
def take_1d_float32_float64(values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=...) -> None: ...
230+
def take_1d_float64_float64(values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=...) -> None: ...
231+
def take_1d_object_object(values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=...) -> None: ...
232+
def take_1d_bool_bool(values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=...) -> None: ...
233+
def take_1d_bool_object(values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=...) -> None: ...
234+
235+
def take_2d_axis0_int8_int8(values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=...) -> None: ...
236+
def take_2d_axis0_int8_int32(values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=...) -> None: ...
237+
def take_2d_axis0_int8_int64(values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=...) -> None: ...
238+
def take_2d_axis0_int8_float64(values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=...) -> None: ...
239+
def take_2d_axis0_int16_int16(values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=...) -> None: ...
240+
def take_2d_axis0_int16_int32(values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=...) -> None: ...
241+
def take_2d_axis0_int16_int64(values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=...) -> None: ...
242+
def take_2d_axis0_int16_float64(values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=...) -> None: ...
243+
def take_2d_axis0_int32_int32(values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=...) -> None: ...
244+
def take_2d_axis0_int32_int64(values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=...) -> None: ...
245+
def take_2d_axis0_int32_float64(values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=...) -> None: ...
246+
def take_2d_axis0_int64_int64(values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=...) -> None: ...
247+
def take_2d_axis0_int64_float64(values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=...) -> None: ...
248+
def take_2d_axis0_float32_float32(values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=...) -> None: ...
249+
def take_2d_axis0_float32_float64(values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=...) -> None: ...
250+
def take_2d_axis0_float64_float64(values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=...) -> None: ...
251+
def take_2d_axis0_object_object(values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=...) -> None: ...
252+
def take_2d_axis0_bool_bool(values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=...) -> None: ...
253+
def take_2d_axis0_bool_object(values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=...) -> None: ...
254+
255+
def take_2d_axis1_int8_int8(values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=...) -> None: ...
256+
def take_2d_axis1_int8_int32(values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=...) -> None: ...
257+
def take_2d_axis1_int8_int64(values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=...) -> None: ...
258+
def take_2d_axis1_int8_float64(values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=...) -> None: ...
259+
def take_2d_axis1_int16_int16(values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=...) -> None: ...
260+
def take_2d_axis1_int16_int32(values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=...) -> None: ...
261+
def take_2d_axis1_int16_int64(values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=...) -> None: ...
262+
def take_2d_axis1_int16_float64(values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=...) -> None: ...
263+
def take_2d_axis1_int32_int32(values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=...) -> None: ...
264+
def take_2d_axis1_int32_int64(values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=...) -> None: ...
265+
def take_2d_axis1_int32_float64(values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=...) -> None: ...
266+
def take_2d_axis1_int64_int64(values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=...) -> None: ...
267+
def take_2d_axis1_int64_float64(values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=...) -> None: ...
268+
def take_2d_axis1_float32_float32(values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=...) -> None: ...
269+
def take_2d_axis1_float32_float64(values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=...) -> None: ...
270+
def take_2d_axis1_float64_float64(values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=...) -> None: ...
271+
def take_2d_axis1_object_object(values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=...) -> None: ...
272+
def take_2d_axis1_bool_bool(values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=...) -> None: ...
273+
def take_2d_axis1_bool_object(values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=...) -> None: ...
274+
275+
def take_2d_multi_int8_int8(values: np.ndarray, indexer, out: np.ndarray, fill_value=...) -> None: ...
276+
def take_2d_multi_int8_int32(values: np.ndarray, indexer, out: np.ndarray, fill_value=...) -> None: ...
277+
def take_2d_multi_int8_int64(values: np.ndarray, indexer, out: np.ndarray, fill_value=...) -> None: ...
278+
def take_2d_multi_int8_float64(values: np.ndarray, indexer, out: np.ndarray, fill_value=...) -> None: ...
279+
def take_2d_multi_int16_int16(values: np.ndarray, indexer, out: np.ndarray, fill_value=...) -> None: ...
280+
def take_2d_multi_int16_int32(values: np.ndarray, indexer, out: np.ndarray, fill_value=...) -> None: ...
281+
def take_2d_multi_int16_int64(values: np.ndarray, indexer, out: np.ndarray, fill_value=...) -> None: ...
282+
def take_2d_multi_int16_float64(values: np.ndarray, indexer, out: np.ndarray, fill_value=...) -> None: ...
283+
def take_2d_multi_int32_int32(values: np.ndarray, indexer, out: np.ndarray, fill_value=...) -> None: ...
284+
def take_2d_multi_int32_int64(values: np.ndarray, indexer, out: np.ndarray, fill_value=...) -> None: ...
285+
def take_2d_multi_int32_float64(values: np.ndarray, indexer, out: np.ndarray, fill_value=...) -> None: ...
286+
def take_2d_multi_int64_float64(values: np.ndarray, indexer, out: np.ndarray, fill_value=...) -> None: ...
287+
def take_2d_multi_float32_float32(values: np.ndarray, indexer, out: np.ndarray, fill_value=...) -> None: ...
288+
def take_2d_multi_float32_float64(values: np.ndarray, indexer, out: np.ndarray, fill_value=...) -> None: ...
289+
def take_2d_multi_float64_float64(values: np.ndarray, indexer, out: np.ndarray, fill_value=...) -> None: ...
290+
def take_2d_multi_object_object(values: np.ndarray, indexer, out: np.ndarray, fill_value=...) -> None: ...
291+
def take_2d_multi_bool_bool(values: np.ndarray, indexer, out: np.ndarray, fill_value=...) -> None: ...
292+
def take_2d_multi_bool_object(values: np.ndarray, indexer, out: np.ndarray, fill_value=...) -> None: ...
293+
def take_2d_multi_int64_int64(values: np.ndarray, indexer, out: np.ndarray, fill_value=...) -> None: ...

‎pandas/_libs/algos.pyx

+1-1
Original file line numberDiff line numberDiff line change
@@ -563,7 +563,7 @@ ctypedef fused algos_t:
563563
uint8_t
564564

565565

566-
def validate_limit(nobs: int, limit=None) -> int:
566+
def validate_limit(nobs: int | None, limit=None) -> int:
567567
"""
568568
Check that the `limit` argument is a positive integer.
569569

‎pandas/core/algorithms.py

+4-13
Original file line numberDiff line numberDiff line change
@@ -155,10 +155,7 @@ def _ensure_data(values: ArrayLike) -> tuple[np.ndarray, DtypeObj]:
155155
with catch_warnings():
156156
simplefilter("ignore", np.ComplexWarning)
157157
values = ensure_float64(values)
158-
# error: Incompatible return value type (got "Tuple[ExtensionArray,
159-
# dtype[floating[_64Bit]]]", expected "Tuple[ndarray, Union[dtype[Any],
160-
# ExtensionDtype]]")
161-
return values, np.dtype("float64") # type: ignore[return-value]
158+
return values, np.dtype("float64")
162159

163160
except (TypeError, ValueError, OverflowError):
164161
# if we are trying to coerce to a dtype
@@ -202,11 +199,7 @@ def _ensure_data(values: ArrayLike) -> tuple[np.ndarray, DtypeObj]:
202199
# we are actually coercing to int64
203200
# until our algos support int* directly (not all do)
204201
values = ensure_int64(values)
205-
206-
# error: Incompatible return value type (got "Tuple[ExtensionArray,
207-
# Union[dtype[Any], ExtensionDtype]]", expected "Tuple[ndarray,
208-
# Union[dtype[Any], ExtensionDtype]]")
209-
return values, dtype # type: ignore[return-value]
202+
return values, dtype
210203

211204
# we have failed, return object
212205
values = np.asarray(values, dtype=object)
@@ -300,7 +293,7 @@ def _get_hashtable_algo(values: np.ndarray):
300293
return htable, values
301294

302295

303-
def _get_values_for_rank(values: ArrayLike):
296+
def _get_values_for_rank(values: ArrayLike) -> np.ndarray:
304297
if is_categorical_dtype(values):
305298
values = cast("Categorical", values)._values_for_rank()
306299

@@ -311,9 +304,7 @@ def _get_values_for_rank(values: ArrayLike):
311304
def get_data_algo(values: ArrayLike):
312305
values = _get_values_for_rank(values)
313306

314-
# error: Argument 1 to "_check_object_for_strings" has incompatible type
315-
# "ExtensionArray"; expected "ndarray"
316-
ndtype = _check_object_for_strings(values) # type: ignore[arg-type]
307+
ndtype = _check_object_for_strings(values)
317308
htable = _hashtables.get(ndtype, _hashtables["object"])
318309

319310
return htable, values

‎pandas/core/array_algos/transforms.py

+5-3
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,6 @@
44

55
import numpy as np
66

7-
from pandas.core.dtypes.common import ensure_platform_int
8-
97

108
def shift(values: np.ndarray, periods: int, axis: int, fill_value) -> np.ndarray:
119
new_values = values
@@ -20,7 +18,11 @@ def shift(values: np.ndarray, periods: int, axis: int, fill_value) -> np.ndarray
2018
axis = new_values.ndim - axis - 1
2119

2220
if new_values.size:
23-
new_values = np.roll(new_values, ensure_platform_int(periods), axis=axis)
21+
new_values = np.roll(
22+
new_values,
23+
np.intp(periods),
24+
axis=axis,
25+
)
2426

2527
axis_indexer = [slice(None)] * values.ndim
2628
if periods > 0:

‎pandas/core/arrays/categorical.py

+6-2
Original file line numberDiff line numberDiff line change
@@ -2576,10 +2576,14 @@ def _get_codes_for_values(values, categories: Index) -> np.ndarray:
25762576
if not isinstance(values, cls):
25772577
# exception raised in _from_sequence
25782578
values = ensure_object(values)
2579-
categories = ensure_object(categories)
2579+
# error: Incompatible types in assignment (expression has type
2580+
# "ndarray", variable has type "Index")
2581+
categories = ensure_object(categories) # type: ignore[assignment]
25802582
elif not dtype_equal:
25812583
values = ensure_object(values)
2582-
categories = ensure_object(categories)
2584+
# error: Incompatible types in assignment (expression has type "ndarray",
2585+
# variable has type "Index")
2586+
categories = ensure_object(categories) # type: ignore[assignment]
25832587

25842588
if isinstance(categories, ABCIndex):
25852589
return coerce_indexer_dtype(categories.get_indexer_for(values), categories)

‎pandas/core/dtypes/cast.py

+1-3
Original file line numberDiff line numberDiff line change
@@ -1226,9 +1226,7 @@ def astype_nansafe(
12261226
from pandas import to_datetime
12271227

12281228
return astype_nansafe(
1229-
# error: No overload variant of "to_datetime" matches argument type
1230-
# "ndarray"
1231-
to_datetime(arr).values, # type: ignore[call-overload]
1229+
to_datetime(arr).values,
12321230
dtype,
12331231
copy=copy,
12341232
)

0 commit comments

Comments
 (0)