Skip to content

Commit 7b22284

Browse files
charrischarris
authored
Merge pull request numpy#19949 from BvB93/rec
ENH: Add annotations for `np.rec`
2 parents 3899d9c + 7e6021e commit 7b22284

File tree

10 files changed

+415
-117
lines changed

10 files changed

+415
-117
lines changed

‎numpy/__init__.pyi‎

Lines changed: 80 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -205,13 +205,13 @@ from numpy import (
205205
matrixlib as matrixlib,
206206
polynomial as polynomial,
207207
random as random,
208-
rec as rec,
209208
testing as testing,
210209
version as version,
211210
)
212211

213-
from numpy.core import defchararray
212+
from numpy.core import defchararray, records
214213
char = defchararray
214+
rec = records
215215

216216
from numpy.core.function_base import (
217217
linspace as linspace,
@@ -706,16 +706,6 @@ class chararray(ndarray[_ShapeType, _DType_co]):
706706
def isnumeric(self): ...
707707
def isdecimal(self): ...
708708

709-
class format_parser:
710-
def __init__(
711-
self,
712-
formats: Any,
713-
names: Any,
714-
titles: Any,
715-
aligned: Any = ...,
716-
byteorder: Any = ...,
717-
) -> None: ...
718-
719709
class matrix(ndarray[_ShapeType, _DType_co]):
720710
def __new__(
721711
subtype,
@@ -804,7 +794,6 @@ class nditer:
804794
def __setitem__(self, index: SupportsIndex | slice, value: Any) -> None: ...
805795
def __delitem__(self, key: SupportsIndex | slice) -> None: ...
806796

807-
808797
class poly1d:
809798
def __init__(
810799
self,
@@ -863,33 +852,6 @@ class poly1d:
863852
def integ(self, m=..., k=...): ...
864853
def deriv(self, m=...): ...
865854

866-
class recarray(ndarray[_ShapeType, _DType_co]):
867-
def __new__(
868-
subtype,
869-
shape: Any,
870-
dtype: Any = ...,
871-
buf: Any = ...,
872-
offset: Any = ...,
873-
strides: Any = ...,
874-
formats: Any = ...,
875-
names: Any = ...,
876-
titles: Any = ...,
877-
byteorder: Any = ...,
878-
aligned: Any = ...,
879-
order: Any = ...,
880-
) -> Any: ...
881-
def __array_finalize__(self, obj): ...
882-
def __getattribute__(self, attr): ...
883-
def __setattr__(self, attr, val): ...
884-
def __getitem__(self, indx): ...
885-
def field(self, attr, val=...): ...
886-
887-
class record(void):
888-
def __getattribute__(self, attr): ...
889-
def __setattr__(self, attr, val): ...
890-
def __getitem__(self, indx): ...
891-
def pprint(self): ...
892-
893855
class vectorize:
894856
pyfunc: Any
895857
cache: Any
@@ -1694,9 +1656,9 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeType, _DType_co]):
16941656
cls: Type[_ArraySelf],
16951657
shape: _ShapeLike,
16961658
dtype: DTypeLike = ...,
1697-
buffer: _SupportsBuffer = ...,
1698-
offset: int = ...,
1699-
strides: _ShapeLike = ...,
1659+
buffer: None | _SupportsBuffer = ...,
1660+
offset: SupportsIndex = ...,
1661+
strides: None | _ShapeLike = ...,
17001662
order: _OrderKACF = ...,
17011663
) -> _ArraySelf: ...
17021664
@overload
@@ -3292,8 +3254,15 @@ class void(flexible):
32923254
def setfield(
32933255
self, val: ArrayLike, dtype: DTypeLike, offset: int = ...
32943256
) -> None: ...
3295-
def __getitem__(self, key: SupportsIndex) -> Any: ...
3296-
def __setitem__(self, key: SupportsIndex, value: ArrayLike) -> None: ...
3257+
@overload
3258+
def __getitem__(self, key: str | SupportsIndex) -> Any: ...
3259+
@overload
3260+
def __getitem__(self, key: list[str]) -> void: ...
3261+
def __setitem__(
3262+
self,
3263+
key: str | List[str] | SupportsIndex,
3264+
value: ArrayLike,
3265+
) -> None: ...
32973266

32983267
void0 = void
32993268

@@ -3704,3 +3673,69 @@ class iinfo(Generic[_IntType]):
37043673
def __new__(cls, dtype: int | Type[int]) -> iinfo[int_]: ...
37053674
@overload
37063675
def __new__(cls, dtype: str) -> iinfo[Any]: ...
3676+
3677+
class format_parser:
3678+
dtype: dtype[void]
3679+
def __init__(
3680+
self,
3681+
formats: DTypeLike,
3682+
names: None | str | Sequence[str],
3683+
titles: None | str | Sequence[str],
3684+
aligned: bool = ...,
3685+
byteorder: None | _ByteOrder = ...,
3686+
) -> None: ...
3687+
3688+
# TODO: field-lookup returns either a `recarray` or a `ndarray`
3689+
# depending on the field dtype
3690+
class recarray(ndarray[_ShapeType, _DType_co]):
3691+
# NOTE: While not strictly mandatory, we're demanding here that arguments
3692+
# for the `format_parser`- and `dtype`-based dtype constructors are
3693+
# mutually exclusive
3694+
@overload
3695+
def __new__(
3696+
subtype,
3697+
shape: _ShapeLike,
3698+
dtype: None = ...,
3699+
buf: None | _SupportsBuffer = ...,
3700+
offset: SupportsIndex = ...,
3701+
strides: None | _ShapeLike = ...,
3702+
*,
3703+
formats: DTypeLike,
3704+
names: None | str | Sequence[str] = ...,
3705+
titles: None | str | Sequence[str] = ...,
3706+
byteorder: None | _ByteOrder = ...,
3707+
aligned: bool = ...,
3708+
order: _OrderKACF = ...,
3709+
) -> recarray[Any, dtype[record]]: ...
3710+
@overload
3711+
def __new__(
3712+
subtype,
3713+
shape: _ShapeLike,
3714+
dtype: DTypeLike,
3715+
buf: None | _SupportsBuffer = ...,
3716+
offset: SupportsIndex = ...,
3717+
strides: None | _ShapeLike = ...,
3718+
formats: None = ...,
3719+
names: None = ...,
3720+
titles: None = ...,
3721+
byteorder: None = ...,
3722+
aligned: L[False] = ...,
3723+
order: _OrderKACF = ...,
3724+
) -> recarray[Any, dtype[Any]]: ...
3725+
def __array_finalize__(self, obj: object) -> None: ...
3726+
def __getattribute__(self, attr: str) -> Any: ...
3727+
def __setattr__(self, attr: str, val: ArrayLike) -> None: ...
3728+
def __getitem__(self, indx): ... # TODO
3729+
@overload
3730+
def field(self, attr: int | str, val: None = ...) -> Any: ...
3731+
@overload
3732+
def field(self, attr: int | str, val: ArrayLike) -> None: ...
3733+
3734+
class record(void):
3735+
def __getattribute__(self, attr: str) -> Any: ...
3736+
def __setattr__(self, attr: str, val: ArrayLike) -> None: ...
3737+
def pprint(self) -> str: ...
3738+
@overload
3739+
def __getitem__(self, key: str | SupportsIndex) -> Any: ...
3740+
@overload
3741+
def __getitem__(self, key: list[str]) -> record: ...

‎numpy/core/records.pyi‎

Lines changed: 183 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,183 @@
1+
import os
2+
from typing import (
3+
List,
4+
Sequence,
5+
Any,
6+
TypeVar,
7+
Iterable,
8+
overload,
9+
Tuple,
10+
Protocol,
11+
)
12+
13+
from numpy import (
14+
format_parser as format_parser,
15+
record as record,
16+
recarray as recarray,
17+
dtype,
18+
generic,
19+
void,
20+
_ByteOrder,
21+
_SupportsBuffer,
22+
)
23+
24+
from numpy.typing import (
25+
ArrayLike,
26+
DTypeLike,
27+
NDArray,
28+
_ShapeLike,
29+
_ArrayLikeVoid_co,
30+
_NestedSequence,
31+
)
32+
33+
_SCT = TypeVar("_SCT", bound=generic)
34+
35+
_RecArray = recarray[Any, dtype[_SCT]]
36+
37+
class _SupportsReadInto(Protocol):
38+
def seek(self, offset: int, whence: int, /) -> object: ...
39+
def tell(self, /) -> int: ...
40+
def readinto(self, buffer: memoryview, /) -> int: ...
41+
42+
__all__: List[str]
43+
44+
@overload
45+
def fromarrays(
46+
arrayList: Iterable[ArrayLike],
47+
dtype: DTypeLike = ...,
48+
shape: None | _ShapeLike = ...,
49+
formats: None = ...,
50+
names: None = ...,
51+
titles: None = ...,
52+
aligned: bool = ...,
53+
byteorder: None = ...,
54+
) -> _RecArray[Any]: ...
55+
@overload
56+
def fromarrays(
57+
arrayList: Iterable[ArrayLike],
58+
dtype: None = ...,
59+
shape: None | _ShapeLike = ...,
60+
*,
61+
formats: DTypeLike,
62+
names: None | str | Sequence[str] = ...,
63+
titles: None | str | Sequence[str] = ...,
64+
aligned: bool = ...,
65+
byteorder: None | _ByteOrder = ...,
66+
) -> _RecArray[record]: ...
67+
68+
@overload
69+
def fromrecords(
70+
recList: _ArrayLikeVoid_co | Tuple[Any, ...] | _NestedSequence[Tuple[Any, ...]],
71+
dtype: DTypeLike = ...,
72+
shape: None | _ShapeLike = ...,
73+
formats: None = ...,
74+
names: None = ...,
75+
titles: None = ...,
76+
aligned: bool = ...,
77+
byteorder: None = ...,
78+
) -> _RecArray[record]: ...
79+
@overload
80+
def fromrecords(
81+
recList: _ArrayLikeVoid_co | Tuple[Any, ...] | _NestedSequence[Tuple[Any, ...]],
82+
dtype: None = ...,
83+
shape: None | _ShapeLike = ...,
84+
*,
85+
formats: DTypeLike,
86+
names: None | str | Sequence[str] = ...,
87+
titles: None | str | Sequence[str] = ...,
88+
aligned: bool = ...,
89+
byteorder: None | _ByteOrder = ...,
90+
) -> _RecArray[record]: ...
91+
92+
@overload
93+
def fromstring(
94+
datastring: _SupportsBuffer,
95+
dtype: DTypeLike,
96+
shape: None | _ShapeLike = ...,
97+
offset: int = ...,
98+
formats: None = ...,
99+
names: None = ...,
100+
titles: None = ...,
101+
aligned: bool = ...,
102+
byteorder: None = ...,
103+
) -> _RecArray[record]: ...
104+
@overload
105+
def fromstring(
106+
datastring: _SupportsBuffer,
107+
dtype: None = ...,
108+
shape: None | _ShapeLike = ...,
109+
offset: int = ...,
110+
*,
111+
formats: DTypeLike,
112+
names: None | str | Sequence[str] = ...,
113+
titles: None | str | Sequence[str] = ...,
114+
aligned: bool = ...,
115+
byteorder: None | _ByteOrder = ...,
116+
) -> _RecArray[record]: ...
117+
118+
@overload
119+
def fromfile(
120+
fd: str | bytes | os.PathLike[str] | os.PathLike[bytes] | _SupportsReadInto,
121+
dtype: DTypeLike,
122+
shape: None | _ShapeLike = ...,
123+
offset: int = ...,
124+
formats: None = ...,
125+
names: None = ...,
126+
titles: None = ...,
127+
aligned: bool = ...,
128+
byteorder: None = ...,
129+
) -> _RecArray[Any]: ...
130+
@overload
131+
def fromfile(
132+
fd: str | bytes | os.PathLike[str] | os.PathLike[bytes] | _SupportsReadInto,
133+
dtype: None = ...,
134+
shape: None | _ShapeLike = ...,
135+
offset: int = ...,
136+
*,
137+
formats: DTypeLike,
138+
names: None | str | Sequence[str] = ...,
139+
titles: None | str | Sequence[str] = ...,
140+
aligned: bool = ...,
141+
byteorder: None | _ByteOrder = ...,
142+
) -> _RecArray[record]: ...
143+
144+
@overload
145+
def array(
146+
obj: _SCT | NDArray[_SCT],
147+
dtype: None = ...,
148+
shape: None | _ShapeLike = ...,
149+
offset: int = ...,
150+
formats: None = ...,
151+
names: None = ...,
152+
titles: None = ...,
153+
aligned: bool = ...,
154+
byteorder: None = ...,
155+
copy: bool = ...,
156+
) -> _RecArray[_SCT]: ...
157+
@overload
158+
def array(
159+
obj: ArrayLike,
160+
dtype: DTypeLike,
161+
shape: None | _ShapeLike = ...,
162+
offset: int = ...,
163+
formats: None = ...,
164+
names: None = ...,
165+
titles: None = ...,
166+
aligned: bool = ...,
167+
byteorder: None = ...,
168+
copy: bool = ...,
169+
) -> _RecArray[Any]: ...
170+
@overload
171+
def array(
172+
obj: ArrayLike,
173+
dtype: None = ...,
174+
shape: None | _ShapeLike = ...,
175+
offset: int = ...,
176+
*,
177+
formats: DTypeLike,
178+
names: None | str | Sequence[str] = ...,
179+
titles: None | str | Sequence[str] = ...,
180+
aligned: bool = ...,
181+
byteorder: None | _ByteOrder = ...,
182+
copy: bool = ...,
183+
) -> _RecArray[record]: ...

‎numpy/lib/npyio.pyi‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ from numpy import (
2929
generic,
3030
float64,
3131
void,
32+
record,
3233
)
3334

3435
from numpy.ma.mrecords import MaskedRecords
@@ -240,7 +241,7 @@ def recfromtxt(
240241
*,
241242
usemask: L[False] = ...,
242243
**kwargs: Any,
243-
) -> recarray[Any, dtype[void]]: ...
244+
) -> recarray[Any, dtype[record]]: ...
244245
@overload
245246
def recfromtxt(
246247
fname: str | os.PathLike[str] | Iterable[str] | Iterable[bytes],
@@ -255,7 +256,7 @@ def recfromcsv(
255256
*,
256257
usemask: L[False] = ...,
257258
**kwargs: Any,
258-
) -> recarray[Any, dtype[void]]: ...
259+
) -> recarray[Any, dtype[record]]: ...
259260
@overload
260261
def recfromcsv(
261262
fname: str | os.PathLike[str] | Iterable[str] | Iterable[bytes],

0 commit comments

Comments
 (0)