22
33from collections .abc import Generator
44from types import EllipsisType
5- from typing import Any , Final , Generic , TypeAlias
5+ from typing import Any , Final , Generic , TypeAlias , overload
66from typing_extensions import TypeVar
77
88import numpy as np
@@ -12,6 +12,7 @@ __all__ = ["Arrayterator"]
1212###
1313
1414_ShapeT_co = TypeVar ("_ShapeT_co" , bound = tuple [int , ...], default = tuple [int , ...], covariant = True )
15+ _DTypeT = TypeVar ("_DTypeT" , bound = np .dtype [Any ])
1516_DTypeT_co = TypeVar ("_DTypeT_co" , bound = np .dtype [Any ], default = np .dtype [Any ], covariant = True )
1617_ScalarT = TypeVar ("_ScalarT" , bound = np .generic )
1718
@@ -37,3 +38,9 @@ class Arrayterator(np.ndarray[_ShapeT_co, _DTypeT_co], Generic[_ShapeT_co, _DTyp
3738 def __init__ (self , / , var : np .ndarray [_ShapeT_co , _DTypeT_co ], buf_size : int | None = None ) -> None : ... # pyright: ignore[reportInconsistentConstructor]
3839 def __getitem__ (self , index : _AnyIndex , / ) -> Arrayterator [tuple [int , ...], _DTypeT_co ]: ... # type: ignore[override]
3940 def __iter__ (self ) -> Generator [np .ndarray [tuple [int , ...], _DTypeT_co ]]: ...
41+
42+ #
43+ @overload # type: ignore[override]
44+ def __array__ (self , / , dtype : _DTypeT_co | None = None , copy : bool | None = None ) -> np .ndarray [_ShapeT_co , _DTypeT_co ]: ...
45+ @overload
46+ def __array__ (self , / , dtype : _DTypeT , copy : bool | None = None ) -> np .ndarray [_ShapeT_co , _DTypeT ]: ...
0 commit comments