Skip to content

Commit f2136a3

Browse files
jbrockmendeljbrockmendel
committed
TYP: lib.pyi
1 parent 5cb0916 commit f2136a3

File tree

18 files changed

+302
-40
lines changed

18 files changed

+302
-40
lines changed

‎pandas/_libs/lib.pyi

+210
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,210 @@
1+
# TODO(npdtypes): Many types specified here can be made more specific/accurate;
2+
# the more specific versions are specified in comments
3+
4+
from typing import (
5+
Any,
6+
Callable,
7+
Union,
8+
)
9+
10+
import numpy as np
11+
12+
from pandas._typing import ArrayLike
13+
14+
# placeholder until we can specify np.ndarray[object, ndim=2]
15+
ndarray_obj_2d = np.ndarray
16+
17+
from enum import Enum
18+
19+
class NoDefault(Enum):
20+
...
21+
22+
no_default: NoDefault
23+
24+
25+
def item_from_zerodim(val: object) -> object: ...
26+
def infer_dtype(value: object, skipna: bool = True) -> str: ...
27+
28+
def is_iterator(obj: object) -> bool: ...
29+
def is_scalar(val: object) -> bool: ...
30+
def is_list_like(obj: object, allow_sets: bool = True) -> bool: ...
31+
32+
def is_period(val: object) -> bool: ...
33+
def is_interval(val: object) -> bool: ...
34+
def is_decimal(val: object) -> bool: ...
35+
def is_complex(val: object) -> bool: ...
36+
def is_bool(val: object) -> bool: ...
37+
def is_integer(val: object) -> bool: ...
38+
def is_float(val: object) -> bool: ...
39+
40+
def is_interval_array(values: np.ndarray) -> bool: ...
41+
def is_period_array(values: np.ndarray) -> bool: ...
42+
def is_datetime64_array(values: np.ndarray) -> bool: ...
43+
def is_timedelta_or_timedelta64_array(values: np.ndarray) -> bool: ...
44+
def is_datetime_with_singletz_array(values: np.ndarray) -> bool: ...
45+
46+
def is_time_array(values: np.ndarray, skipna: bool = False): ...
47+
def is_date_array(values: np.ndarray, skipna: bool = False): ...
48+
def is_datetime_array(values: np.ndarray, skipna: bool = False): ...
49+
def is_string_array(values: np.ndarray, skipna: bool = False): ...
50+
def is_float_array(values: np.ndarray, skipna: bool = False): ...
51+
def is_integer_array(values: np.ndarray, skipna: bool = False): ...
52+
def is_bool_array(values: np.ndarray, skipna: bool = False): ...
53+
54+
def fast_multiget(mapping: dict, keys: np.ndarray, default=np.nan) -> ArrayLike: ...
55+
56+
# TODO: gen: Generator?
57+
def fast_unique_multiple_list_gen(gen: object, sort: bool = True) -> list: ...
58+
def fast_unique_multiple_list(lists: list, sort: bool = True) -> list: ...
59+
def fast_unique_multiple(arrays: list, sort: bool = True) -> list: ...
60+
61+
def map_infer(
62+
arr: np.ndarray, f: Callable[[Any], Any], convert: bool = True, ignore_na: bool = False
63+
) -> ArrayLike: ...
64+
65+
def maybe_convert_objects(
66+
objects: np.ndarray, # np.ndarray[object]
67+
try_float: bool = False,
68+
safe: bool = False,
69+
convert_datetime: bool = False,
70+
convert_timedelta: bool = False,
71+
convert_to_nullable_integer: bool = False,
72+
) -> ArrayLike: ...
73+
74+
def maybe_convert_numeric(
75+
values: np.ndarray, # np.ndarray[object]
76+
na_values: set,
77+
convert_empty: bool = True,
78+
coerce_numeric: bool = False,
79+
) -> np.ndarray: ...
80+
81+
# TODO: restrict `arr`?
82+
def ensure_string_array(
83+
arr,
84+
na_value: object = np.nan,
85+
convert_na_value: bool = True,
86+
copy: bool = True,
87+
skipna: bool = True,
88+
) -> np.ndarray: ... # np.ndarray[object]
89+
90+
def infer_datetimelike_array(
91+
arr: np.ndarray # np.ndarray[object]
92+
) -> str: ...
93+
94+
# TODO: new_dtype -> np.dtype?
95+
def astype_intsafe(
96+
arr: np.ndarray, # np.ndarray[object]
97+
new_dtype,
98+
) -> np.ndarray: ...
99+
100+
def fast_zip(ndarrays: list) -> np.ndarray: ... # np.ndarray[object]
101+
102+
# TODO: can we be more specific about rows?
103+
def to_object_array_tuples(rows: object) -> ndarray_obj_2d: ...
104+
105+
def tuples_to_object_array(
106+
tuples: np.ndarray # np.ndarray[object]
107+
) -> ndarray_obj_2d: ...
108+
109+
# TODO: can we be more specific about rows?
110+
def to_object_array(rows: object, min_width: int = 0) -> ndarray_obj_2d: ...
111+
112+
def dicts_to_array(dicts: list, columns: list) -> ndarray_obj_2d: ...
113+
114+
115+
def maybe_booleans_to_slice(
116+
mask: np.ndarray # ndarray[uint8_t]
117+
) -> Union[
118+
slice,
119+
np.ndarray, # np.ndarray[np.uint8]
120+
]: ...
121+
122+
def maybe_indices_to_slice(
123+
indices: np.ndarray, # np.ndarray[np.intp]
124+
max_len: int,
125+
) -> Union[
126+
slice,
127+
np.ndarray, # np.ndarray[np.intp]
128+
]: ...
129+
130+
def clean_index_list(obj: list) -> tuple[
131+
Union[
132+
list,
133+
np.ndarray, # np.ndarray[object] | np.ndarray[np.int64]
134+
],
135+
bool,
136+
]: ...
137+
138+
139+
# -----------------------------------------------------------------
140+
# Functions which in reality take memoryviews
141+
142+
def memory_usage_of_objects(
143+
arr: np.ndarray # object[:]
144+
) -> int: ... # np.int64
145+
146+
147+
# TODO: f: Callable?
148+
# TODO: dtype -> DtypeObj?
149+
def map_infer_mask(
150+
arr: np.ndarray,
151+
f: Callable[[Any], Any],
152+
mask: np.ndarray, # const uint8_t[:]
153+
convert: bool = ...,
154+
na_value: Any = ...,
155+
dtype: Any = ...,
156+
) -> ArrayLike: ...
157+
158+
def indices_fast(
159+
index: np.ndarray, # ndarray[intp_t]
160+
labels: np.ndarray, # const int64_t[:]
161+
keys: list,
162+
sorted_labels: list[np.ndarray], # list[ndarray[np.int64]]
163+
) -> dict: ...
164+
165+
def generate_slices(
166+
labels: np.ndarray, # const intp_t[:]
167+
ngroups: int
168+
) -> tuple[
169+
np.ndarray, # np.ndarray[np.int64]
170+
np.ndarray, # np.ndarray[np.int64]
171+
]: ...
172+
173+
def count_level_2d(
174+
mask: np.ndarray, # ndarray[uint8_t, ndim=2, cast=True],
175+
labels: np.ndarray, # const intp_t[:]
176+
max_bin: int,
177+
axis: int
178+
) -> np.ndarray: ... # np.ndarray[np.int64, ndim=2]
179+
180+
def get_level_sorter(
181+
label: np.ndarray, # const int64_t[:]
182+
starts: np.ndarray, # const intp_t[:]
183+
) -> np.ndarray: ... # np.ndarray[np.intp, ndim=1]
184+
185+
186+
def generate_bins_dt64(
187+
values: np.ndarray, # np.ndarray[np.int64]
188+
binner: np.ndarray, # const int64_t[:]
189+
closed: object = "left",
190+
hasnans: bool = False,
191+
) -> np.ndarray: ... # np.ndarray[np.int64, ndim=1]
192+
193+
194+
def array_equivalent_object(
195+
left: np.ndarray, # object[:]
196+
right: np.ndarray, # object[:]
197+
) -> bool: ...
198+
199+
def has_infs_f8(
200+
arr: np.ndarray # const float64_t[:]
201+
) -> bool: ...
202+
203+
def has_infs_f4(
204+
arr: np.ndarray # const float32_t[:]
205+
) -> bool: ...
206+
207+
def get_reverse_indexer(
208+
indexer: np.ndarray, # const intp_t[:]
209+
length: int,
210+
) -> np.ndarray: ... # np.ndarray[np.intp]

‎pandas/_testing/asserters.py

+16-5
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,10 @@
66

77
import numpy as np
88

9-
from pandas._libs.lib import no_default
9+
from pandas._libs.lib import (
10+
NoDefault,
11+
no_default,
12+
)
1013
from pandas._libs.missing import is_matching_na
1114
import pandas._libs.testing as _testing
1215

@@ -54,7 +57,7 @@ def assert_almost_equal(
5457
left,
5558
right,
5659
check_dtype: Union[bool, str] = "equiv",
57-
check_less_precise: Union[bool, int] = no_default,
60+
check_less_precise: Union[bool, int, NoDefault] = no_default,
5861
rtol: float = 1.0e-5,
5962
atol: float = 1.0e-8,
6063
**kwargs,
@@ -104,7 +107,11 @@ def assert_almost_equal(
104107
FutureWarning,
105108
stacklevel=2,
106109
)
107-
rtol = atol = _get_tol_from_less_precise(check_less_precise)
110+
# error: Argument 1 to "_get_tol_from_less_precise" has incompatible
111+
# type "Union[bool, int, NoDefault]"; expected "Union[bool, int]"
112+
rtol = atol = _get_tol_from_less_precise(
113+
check_less_precise # type: ignore[arg-type]
114+
)
108115

109116
if isinstance(left, Index):
110117
assert_index_equal(
@@ -242,7 +249,7 @@ def assert_index_equal(
242249
right: Index,
243250
exact: Union[bool, str] = "equiv",
244251
check_names: bool = True,
245-
check_less_precise: Union[bool, int] = no_default,
252+
check_less_precise: Union[bool, int, NoDefault] = no_default,
246253
check_exact: bool = True,
247254
check_categorical: bool = True,
248255
check_order: bool = True,
@@ -331,7 +338,11 @@ def _get_ilevel_values(index, level):
331338
FutureWarning,
332339
stacklevel=2,
333340
)
334-
rtol = atol = _get_tol_from_less_precise(check_less_precise)
341+
# error: Argument 1 to "_get_tol_from_less_precise" has incompatible
342+
# type "Union[bool, int, NoDefault]"; expected "Union[bool, int]"
343+
rtol = atol = _get_tol_from_less_precise(
344+
check_less_precise # type: ignore[arg-type]
345+
)
335346

336347
# instance validation
337348
_check_isinstance(left, right, Index)

‎pandas/core/apply.py

+10-1
Original file line numberDiff line numberDiff line change
@@ -1026,7 +1026,16 @@ def apply_standard(self) -> FrameOrSeriesUnion:
10261026
mapped = obj._values.map(f)
10271027
else:
10281028
values = obj.astype(object)._values
1029-
mapped = lib.map_infer(values, f, convert=self.convert_dtype)
1029+
# error: Argument 2 to "map_infer" has incompatible type
1030+
# "Union[Callable[..., Any], str, List[Union[Callable[..., Any], str]],
1031+
# Dict[Hashable, Union[Union[Callable[..., Any], str],
1032+
# List[Union[Callable[..., Any], str]]]]]"; expected
1033+
# "Callable[[Any], Any]"
1034+
mapped = lib.map_infer(
1035+
values,
1036+
f, # type: ignore[arg-type]
1037+
convert=self.convert_dtype,
1038+
)
10301039

10311040
if len(mapped) and isinstance(mapped[0], ABCSeries):
10321041
# GH 25959 use pd.array instead of tolist

‎pandas/core/arrays/datetimelike.py

+3-1
Original file line numberDiff line numberDiff line change
@@ -293,7 +293,9 @@ def _box_values(self, values) -> np.ndarray:
293293
"""
294294
apply box func to passed values
295295
"""
296-
return lib.map_infer(values, self._box_func)
296+
# error: Incompatible return value type (got
297+
# "Union[ExtensionArray, ndarray]", expected "ndarray")
298+
return lib.map_infer(values, self._box_func) # type: ignore[return-value]
297299

298300
def __iter__(self):
299301
if self.ndim > 1:

‎pandas/core/arrays/integer.py

+3-3
Original file line numberDiff line numberDiff line change
@@ -371,14 +371,14 @@ def astype(self, dtype, copy: bool = True) -> ArrayLike:
371371
if isinstance(dtype, ExtensionDtype):
372372
return super().astype(dtype, copy=copy)
373373

374+
na_value: float | np.datetime64 | lib.NoDefault
375+
374376
# coerce
375377
if is_float_dtype(dtype):
376378
# In astype, we consider dtype=float to also mean na_value=np.nan
377379
na_value = np.nan
378380
elif is_datetime64_dtype(dtype):
379-
# error: Incompatible types in assignment (expression has type
380-
# "datetime64", variable has type "float")
381-
na_value = np.datetime64("NaT") # type: ignore[assignment]
381+
na_value = np.datetime64("NaT")
382382
else:
383383
na_value = lib.no_default
384384

‎pandas/core/arrays/string_.py

+3-1
Original file line numberDiff line numberDiff line change
@@ -445,7 +445,9 @@ def _str_map(self, f, na_value=None, dtype: Optional[Dtype] = None):
445445
if not na_value_is_na:
446446
mask[:] = False
447447

448-
return constructor(result, mask)
448+
# error: Argument 1 to "maybe_convert_objects" has incompatible
449+
# type "Union[ExtensionArray, ndarray]"; expected "ndarray"
450+
return constructor(result, mask) # type: ignore[arg-type]
449451

450452
elif is_string_dtype(dtype) and not is_object_dtype(dtype):
451453
# i.e. StringDtype

‎pandas/core/base.py

+2-1
Original file line numberDiff line numberDiff line change
@@ -1194,7 +1194,8 @@ def _memory_usage(self, deep: bool = False) -> int:
11941194

11951195
v = self.array.nbytes
11961196
if deep and is_object_dtype(self) and not PYPY:
1197-
v += lib.memory_usage_of_objects(self._values)
1197+
values = cast(np.ndarray, self._values)
1198+
v += lib.memory_usage_of_objects(values)
11981199
return v
11991200

12001201
@doc(

‎pandas/core/dtypes/cast.py

+14-8
Original file line numberDiff line numberDiff line change
@@ -128,18 +128,19 @@ def maybe_convert_platform(
128128
values: Union[list, tuple, range, np.ndarray, ExtensionArray]
129129
) -> ArrayLike:
130130
""" try to do platform conversion, allow ndarray or list here """
131+
arr: ArrayLike
132+
131133
if isinstance(values, (list, tuple, range)):
132134
arr = construct_1d_object_array_from_listlike(values)
133135
else:
134136
# The caller is responsible for ensuring that we have np.ndarray
135137
# or ExtensionArray here.
136-
137-
# error: Incompatible types in assignment (expression has type "Union[ndarray,
138-
# ExtensionArray]", variable has type "ndarray")
139-
arr = values # type: ignore[assignment]
138+
arr = values
140139

141140
if arr.dtype == object:
142-
arr = lib.maybe_convert_objects(arr)
141+
# error: Argument 1 to "maybe_convert_objects" has incompatible type
142+
# "Union[ExtensionArray, ndarray]"; expected "ndarray"
143+
arr = lib.maybe_convert_objects(arr) # type: ignore[arg-type]
143144

144145
return arr
145146

@@ -1400,11 +1401,13 @@ def soft_convert_objects(
14001401
# GH 20380, when datetime is beyond year 2262, hence outside
14011402
# bound of nanosecond-resolution 64-bit integers.
14021403
try:
1403-
values = lib.maybe_convert_objects(
1404+
converted = lib.maybe_convert_objects(
14041405
values, convert_datetime=datetime, convert_timedelta=timedelta
14051406
)
14061407
except (OutOfBoundsDatetime, ValueError):
14071408
return values
1409+
if converted is not values:
1410+
return converted
14081411

14091412
if numeric and is_object_dtype(values.dtype):
14101413
converted = lib.maybe_convert_numeric(values, set(), coerce_numeric=True)
@@ -1443,13 +1446,16 @@ def convert_dtypes(
14431446
14441447
Returns
14451448
-------
1449+
str, np.dtype, or ExtensionDtype
14461450
dtype
14471451
new dtype
14481452
"""
1449-
is_extension = is_extension_array_dtype(input_array.dtype)
1453+
inferred_dtype: str | np.dtype | ExtensionDtype
1454+
# TODO: rule out str
1455+
14501456
if (
14511457
convert_string or convert_integer or convert_boolean or convert_floating
1452-
) and not is_extension:
1458+
) and isinstance(input_array, np.ndarray):
14531459
inferred_dtype = lib.infer_dtype(input_array)
14541460

14551461
if not convert_string and is_string_dtype(inferred_dtype):

0 commit comments

Comments
 (0)