diff options
Diffstat (limited to 'numpy')
| -rw-r--r-- | numpy/__init__.pyi | 14 | ||||
| -rw-r--r-- | numpy/typing/tests/data/reveal/flatiter.pyi | 6 |
2 files changed, 18 insertions, 2 deletions
diff --git a/numpy/__init__.pyi b/numpy/__init__.pyi index 062df938d..e18d074de 100644 --- a/numpy/__init__.pyi +++ b/numpy/__init__.pyi @@ -944,12 +944,22 @@ class flatiter(Generic[_NdArraySubClass]): @overload def __getitem__( self: flatiter[ndarray[Any, dtype[_ScalarType]]], - key: int | integer, + key: int | integer | tuple[int | integer], ) -> _ScalarType: ... @overload def __getitem__( - self, key: _ArrayLikeInt | slice | ellipsis + self, + key: _ArrayLikeInt | slice | ellipsis | tuple[_ArrayLikeInt | slice | ellipsis], ) -> _NdArraySubClass: ... + # TODO: `__setitem__` operates via `unsafe` casting rules, and can + # thus accept any type accepted by the relevant underlying `np.generic` + # constructor. + # This means that `value` must in reality be a supertype of `npt.ArrayLike`. + def __setitem__( + self, + key: _ArrayLikeInt | slice | ellipsis | tuple[_ArrayLikeInt | slice | ellipsis], + value: Any, + ) -> None: ... @overload def __array__(self: flatiter[ndarray[Any, _DType]], dtype: None = ..., /) -> ndarray[Any, _DType]: ... @overload diff --git a/numpy/typing/tests/data/reveal/flatiter.pyi b/numpy/typing/tests/data/reveal/flatiter.pyi index ef89acb58..9e7d6e545 100644 --- a/numpy/typing/tests/data/reveal/flatiter.pyi +++ b/numpy/typing/tests/data/reveal/flatiter.pyi @@ -13,5 +13,11 @@ reveal_type(a[0]) # E: str_ reveal_type(a[[0, 1, 2]]) # E: ndarray[Any, dtype[str_]] reveal_type(a[...]) # E: ndarray[Any, dtype[str_]] reveal_type(a[:]) # E: ndarray[Any, dtype[str_]] +reveal_type(a[(...,)]) # E: ndarray[Any, dtype[str_]] +reveal_type(a[(0,)]) # E: str_ reveal_type(a.__array__()) # E: ndarray[Any, dtype[str_]] reveal_type(a.__array__(np.dtype(np.float64))) # E: ndarray[Any, dtype[{float64}]] +a[0] = "a" +a[:5] = "a" +a[...] = "a" +a[(...,)] = "a" |