Skip to content

Commit 26e656e

Browse files
charrischarris
authored
Merge pull request numpy#19953 from BvB93/nditer
ENH: Add annotations to `np.core.multiarray` part 4/4
2 parents 0b01b48 + 2fa73cd commit 26e656e

File tree

7 files changed

+189
-48
lines changed

7 files changed

+189
-48
lines changed

‎numpy/__init__.pyi‎

Lines changed: 110 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,6 @@ if sys.version_info >= (3, 9):
1313
from types import GenericAlias
1414

1515
from numpy._pytesttester import PytestTester
16-
from numpy.core.multiarray import flagsobj
1716
from numpy.core._internal import _ctypes
1817
from numpy.core.getlimits import MachArLike
1918

@@ -196,6 +195,7 @@ from typing import (
196195
Protocol,
197196
SupportsIndex,
198197
Final,
198+
final,
199199
)
200200

201201
# Ensures that the stubs are picked up
@@ -351,6 +351,8 @@ from numpy.core.multiarray import (
351351
geterrobj as geterrobj,
352352
fromstring as fromstring,
353353
frompyfunc as frompyfunc,
354+
nested_iters as nested_iters,
355+
flagsobj,
354356
)
355357

356358
from numpy.core.numeric import (
@@ -768,35 +770,6 @@ class memmap(ndarray[_ShapeType, _DType_co]):
768770
) -> Any: ...
769771
def __getattr__(self, key: str) -> Any: ...
770772

771-
class nditer:
772-
def __new__(
773-
cls,
774-
op: Any,
775-
flags: Any = ...,
776-
op_flags: Any = ...,
777-
op_dtypes: Any = ...,
778-
order: Any = ...,
779-
casting: Any = ...,
780-
op_axes: Any = ...,
781-
itershape: Any = ...,
782-
buffersize: Any = ...,
783-
) -> Any: ...
784-
def __getattr__(self, key: str) -> Any: ...
785-
def __enter__(self) -> nditer: ...
786-
def __exit__(
787-
self,
788-
exc_type: None | Type[BaseException],
789-
exc_value: None | BaseException,
790-
traceback: None | TracebackType,
791-
) -> None: ...
792-
def __iter__(self) -> Iterator[Any]: ...
793-
def __next__(self) -> Any: ...
794-
def __len__(self) -> int: ...
795-
def __copy__(self) -> nditer: ...
796-
def __getitem__(self, index: SupportsIndex | slice) -> Any: ...
797-
def __setitem__(self, index: SupportsIndex | slice, value: Any) -> None: ...
798-
def __delitem__(self, key: SupportsIndex | slice) -> None: ...
799-
800773
class poly1d:
801774
def __init__(
802775
self,
@@ -885,9 +858,6 @@ alltrue = all
885858

886859
def show_config() -> None: ...
887860

888-
# TODO: Sort out which parameters are positional-only
889-
def nested_iters(*args, **kwargs): ... # TODO: Sort out parameters
890-
891861
_NdArraySubClass = TypeVar("_NdArraySubClass", bound=ndarray)
892862
_DTypeScalar_co = TypeVar("_DTypeScalar_co", covariant=True, bound=generic)
893863
_ByteOrder = L["S", "<", ">", "=", "|", "L", "B", "N", "I"]
@@ -3751,3 +3721,110 @@ class record(void):
37513721
def __getitem__(self, key: str | SupportsIndex) -> Any: ...
37523722
@overload
37533723
def __getitem__(self, key: list[str]) -> record: ...
3724+
3725+
_NDIterFlagsKind = L[
3726+
"buffered",
3727+
"c_index",
3728+
"copy_if_overlap",
3729+
"common_dtype",
3730+
"delay_bufalloc",
3731+
"external_loop",
3732+
"f_index",
3733+
"grow_inner", "growinner",
3734+
"multi_index",
3735+
"ranged",
3736+
"refs_ok",
3737+
"reduce_ok",
3738+
"zerosize_ok",
3739+
]
3740+
3741+
_NDIterOpFlagsKind = L[
3742+
"aligned",
3743+
"allocate",
3744+
"arraymask",
3745+
"copy",
3746+
"config",
3747+
"nbo",
3748+
"no_subtype",
3749+
"no_broadcast",
3750+
"overlap_assume_elementwise",
3751+
"readonly",
3752+
"readwrite",
3753+
"updateifcopy",
3754+
"virtual",
3755+
"writeonly",
3756+
"writemasked"
3757+
]
3758+
3759+
@final
3760+
class nditer:
3761+
def __new__(
3762+
cls,
3763+
op: ArrayLike | Sequence[ArrayLike],
3764+
flags: None | Sequence[_NDIterFlagsKind] = ...,
3765+
op_flags: None | Sequence[Sequence[_NDIterOpFlagsKind]] = ...,
3766+
op_dtypes: DTypeLike | Sequence[DTypeLike] = ...,
3767+
order: _OrderKACF = ...,
3768+
casting: _CastingKind = ...,
3769+
op_axes: None | Sequence[Sequence[SupportsIndex]] = ...,
3770+
itershape: None | _ShapeLike = ...,
3771+
buffersize: SupportsIndex = ...,
3772+
) -> nditer: ...
3773+
def __enter__(self) -> nditer: ...
3774+
def __exit__(
3775+
self,
3776+
exc_type: None | Type[BaseException],
3777+
exc_value: None | BaseException,
3778+
traceback: None | TracebackType,
3779+
) -> None: ...
3780+
def __iter__(self) -> nditer: ...
3781+
def __next__(self) -> Tuple[NDArray[Any], ...]: ...
3782+
def __len__(self) -> int: ...
3783+
def __copy__(self) -> nditer: ...
3784+
@overload
3785+
def __getitem__(self, index: SupportsIndex) -> NDArray[Any]: ...
3786+
@overload
3787+
def __getitem__(self, index: slice) -> Tuple[NDArray[Any], ...]: ...
3788+
def __setitem__(self, index: slice | SupportsIndex, value: ArrayLike) -> None: ...
3789+
def close(self) -> None: ...
3790+
def copy(self) -> nditer: ...
3791+
def debug_print(self) -> None: ...
3792+
def enable_external_loop(self) -> None: ...
3793+
def iternext(self) -> bool: ...
3794+
def remove_axis(self, i: SupportsIndex, /) -> None: ...
3795+
def remove_multi_index(self) -> None: ...
3796+
def reset(self) -> None: ...
3797+
@property
3798+
def dtypes(self) -> Tuple[dtype[Any], ...]: ...
3799+
@property
3800+
def finished(self) -> bool: ...
3801+
@property
3802+
def has_delayed_bufalloc(self) -> bool: ...
3803+
@property
3804+
def has_index(self) -> bool: ...
3805+
@property
3806+
def has_multi_index(self) -> bool: ...
3807+
@property
3808+
def index(self) -> int: ...
3809+
@property
3810+
def iterationneedsapi(self) -> bool: ...
3811+
@property
3812+
def iterindex(self) -> int: ...
3813+
@property
3814+
def iterrange(self) -> Tuple[int, ...]: ...
3815+
@property
3816+
def itersize(self) -> int: ...
3817+
@property
3818+
def itviews(self) -> Tuple[NDArray[Any], ...]: ...
3819+
@property
3820+
def multi_index(self) -> Tuple[int, ...]: ...
3821+
@property
3822+
def ndim(self) -> int: ...
3823+
@property
3824+
def nop(self) -> int: ...
3825+
@property
3826+
def operands(self) -> Tuple[NDArray[Any], ...]: ...
3827+
@property
3828+
def shape(self) -> Tuple[int, ...]: ...
3829+
@property
3830+
def value(self) -> Tuple[NDArray[Any], ...]: ...

‎numpy/core/_add_newdocs.py‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -478,7 +478,7 @@
478478

479479
add_newdoc('numpy.core', 'nditer', ('remove_axis',
480480
"""
481-
remove_axis(i)
481+
remove_axis(i, /)
482482
483483
Removes axis `i` from the iterator. Requires that the flag "multi_index"
484484
be enabled.
@@ -504,6 +504,9 @@
504504

505505
add_newdoc('numpy.core', 'nested_iters',
506506
"""
507+
nested_iters(op, axes, flags=None, op_flags=None, op_dtypes=None, \
508+
order="K", casting="safe", buffersize=0)
509+
507510
Create nditers for use in nested loops
508511
509512
Create a tuple of `nditer` objects which iterate in nested loops over

‎numpy/core/multiarray.pyi‎

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,6 @@ from numpy import (
3030
nditer as nditer,
3131

3232
# The rest
33-
nditer,
3433
ufunc,
3534
str_,
3635
bool_,
@@ -51,6 +50,8 @@ from numpy import (
5150
_ModeKind,
5251
_SupportsBuffer,
5352
_IOProtocol,
53+
_NDIterFlagsKind,
54+
_NDIterOpFlagsKind,
5455
)
5556

5657
from numpy.typing import (
@@ -1012,3 +1013,14 @@ class flagsobj:
10121013
def owndata(self) -> bool: ...
10131014
def __getitem__(self, key: _GetItemKeys) -> bool: ...
10141015
def __setitem__(self, key: _SetItemKeys, value: bool) -> None: ...
1016+
1017+
def nested_iters(
1018+
op: ArrayLike | Sequence[ArrayLike],
1019+
axes: Sequence[Sequence[SupportsIndex]],
1020+
flags: None | Sequence[_NDIterFlagsKind] = ...,
1021+
op_flags: None | Sequence[Sequence[_NDIterOpFlagsKind]] = ...,
1022+
op_dtypes: DTypeLike | Sequence[DTypeLike] = ...,
1023+
order: _OrderKACF = ...,
1024+
casting: _CastingKind = ...,
1025+
buffersize: SupportsIndex = ...,
1026+
) -> Tuple[nditer, ...]: ...

‎numpy/typing/tests/data/fail/multiarray.py‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,3 +47,10 @@ def func(a: int) -> None: ...
4747
np.compare_chararrays("a", b"a", "==", False) # E: No overload variant
4848

4949
np.add_docstring(func, None) # E: incompatible type
50+
51+
np.nested_iters([AR_i8, AR_i8]) # E: Missing positional argument
52+
np.nested_iters([AR_i8, AR_i8], 0) # E: incompatible type
53+
np.nested_iters([AR_i8, AR_i8], [0]) # E: incompatible type
54+
np.nested_iters([AR_i8, AR_i8], [[0], [1]], flags=["test"]) # E: incompatible type
55+
np.nested_iters([AR_i8, AR_i8], [[0], [1]], op_flags=[["test"]]) # E: incompatible type
56+
np.nested_iters([AR_i8, AR_i8], [[0], [1]], buffersize=1.0) # E: incompatible type

‎numpy/typing/tests/data/fail/nditer.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
import numpy as np
2+
3+
class Test(np.nditer): ... # E: Cannot inherit from final class
4+
5+
np.nditer([0, 1], flags=["test"]) # E: incompatible type
6+
np.nditer([0, 1], op_flags=[["test"]]) # E: incompatible type
7+
np.nditer([0, 1], itershape=(1.0,)) # E: incompatible type
8+
np.nditer([0, 1], buffersize=1.0) # E: incompatible type

‎numpy/typing/tests/data/reveal/multiarray.py‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,8 @@ class SubClass(np.ndarray[Any, np.dtype[_SCT]]): ...
2525
b_f8 = np.broadcast(AR_f8)
2626
b_i8_f8_f8 = np.broadcast(AR_i8, AR_f8, AR_f8)
2727

28+
nditer_obj: np.nditer
29+
2830
def func(a: int) -> bool: ...
2931

3032
reveal_type(next(b_f8)) # E: tuple[Any]
@@ -123,3 +125,8 @@ def func(a: int) -> bool: ...
123125
reveal_type(np.compare_chararrays(b"a", b"a", "==", True)) # E: numpy.ndarray[Any, numpy.dtype[numpy.bool_]]
124126

125127
reveal_type(np.add_docstring(func, "test")) # E: None
128+
129+
reveal_type(np.nested_iters([AR_i8, AR_i8], [[0], [1]], flags=["c_index"])) # E: tuple[numpy.nditer]
130+
reveal_type(np.nested_iters([AR_i8, AR_i8], [[0], [1]], op_flags=[["readonly", "readonly"]])) # E: tuple[numpy.nditer]
131+
reveal_type(np.nested_iters([AR_i8, AR_i8], [[0], [1]], op_dtypes=np.int_)) # E: tuple[numpy.nditer]
132+
reveal_type(np.nested_iters([AR_i8, AR_i8], [[0], [1]], order="C", casting="no")) # E: tuple[numpy.nditer]

‎numpy/typing/tests/data/reveal/nditer.py‎

Lines changed: 40 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,46 @@
1-
import copy
21
import numpy as np
32

43
nditer_obj: np.nditer
54

6-
with nditer_obj as context:
7-
reveal_type(context) # E: numpy.nditer
5+
reveal_type(np.nditer([0, 1], flags=["c_index"])) # E: numpy.nditer
6+
reveal_type(np.nditer([0, 1], op_flags=[["readonly", "readonly"]])) # E: numpy.nditer
7+
reveal_type(np.nditer([0, 1], op_dtypes=np.int_)) # E: numpy.nditer
8+
reveal_type(np.nditer([0, 1], order="C", casting="no")) # E: numpy.nditer
89

9-
reveal_type(len(nditer_obj)) # E: builtins.int
10-
reveal_type(copy.copy(nditer_obj)) # E: numpy.nditer
11-
reveal_type(next(nditer_obj)) # E: Any
12-
reveal_type(iter(nditer_obj)) # E: typing.Iterator[Any]
13-
reveal_type(nditer_obj[1]) # E: Any
14-
reveal_type(nditer_obj[1:5]) # E: Any
10+
reveal_type(nditer_obj.dtypes) # E: tuple[numpy.dtype[Any]]
11+
reveal_type(nditer_obj.finished) # E: bool
12+
reveal_type(nditer_obj.has_delayed_bufalloc) # E: bool
13+
reveal_type(nditer_obj.has_index) # E: bool
14+
reveal_type(nditer_obj.has_multi_index) # E: bool
15+
reveal_type(nditer_obj.index) # E: int
16+
reveal_type(nditer_obj.iterationneedsapi) # E: bool
17+
reveal_type(nditer_obj.iterindex) # E: int
18+
reveal_type(nditer_obj.iterrange) # E: tuple[builtins.int]
19+
reveal_type(nditer_obj.itersize) # E: int
20+
reveal_type(nditer_obj.itviews) # E: tuple[numpy.ndarray[Any, numpy.dtype[Any]]]
21+
reveal_type(nditer_obj.multi_index) # E: tuple[builtins.int]
22+
reveal_type(nditer_obj.ndim) # E: int
23+
reveal_type(nditer_obj.nop) # E: int
24+
reveal_type(nditer_obj.operands) # E: tuple[numpy.ndarray[Any, numpy.dtype[Any]]]
25+
reveal_type(nditer_obj.shape) # E: tuple[builtins.int]
26+
reveal_type(nditer_obj.value) # E: tuple[numpy.ndarray[Any, numpy.dtype[Any]]]
1527

16-
nditer_obj[1] = 1
17-
nditer_obj[1:5] = 1
18-
del nditer_obj[1]
19-
del nditer_obj[1:5]
28+
reveal_type(nditer_obj.close()) # E: None
29+
reveal_type(nditer_obj.copy()) # E: numpy.nditer
30+
reveal_type(nditer_obj.debug_print()) # E: None
31+
reveal_type(nditer_obj.enable_external_loop()) # E: None
32+
reveal_type(nditer_obj.iternext()) # E: bool
33+
reveal_type(nditer_obj.remove_axis(0)) # E: None
34+
reveal_type(nditer_obj.remove_multi_index()) # E: None
35+
reveal_type(nditer_obj.reset()) # E: None
36+
37+
reveal_type(len(nditer_obj)) # E: int
38+
reveal_type(iter(nditer_obj)) # E: Iterator[builtins.tuple[numpy.ndarray[Any, numpy.dtype[Any]]]]
39+
reveal_type(next(nditer_obj)) # E: tuple[numpy.ndarray[Any, numpy.dtype[Any]]]
40+
reveal_type(nditer_obj.__copy__()) # E: numpy.nditer
41+
with nditer_obj as f:
42+
reveal_type(f) # E: numpy.nditer
43+
reveal_type(nditer_obj[0]) # E: numpy.ndarray[Any, numpy.dtype[Any]]
44+
reveal_type(nditer_obj[:]) # E: tuple[numpy.ndarray[Any, numpy.dtype[Any]]]
45+
nditer_obj[0] = 0
46+
nditer_obj[:] = [0, 1]

0 commit comments

Comments
 (0)