diff options
| -rw-r--r-- | numpy/__init__.pyi | 81 |
1 files changed, 71 insertions, 10 deletions
diff --git a/numpy/__init__.pyi b/numpy/__init__.pyi index 4c78179de..7abf9c62e 100644 --- a/numpy/__init__.pyi +++ b/numpy/__init__.pyi @@ -884,6 +884,13 @@ class dtype(Generic[_DTypeScalar_co]): def __ge__(self, other: DTypeLike) -> bool: ... def __lt__(self, other: DTypeLike) -> bool: ... def __le__(self, other: DTypeLike) -> bool: ... + + # Explicitly defined `__eq__` and `__ne__` to get around mypy's + # `strict_equality` option; even though their signatures are + # identical to their `object`-based counterpart + def __eq__(self, other: Any) -> bool: ... + def __ne__(self, other: Any) -> bool: ... + @property def alignment(self) -> int: ... @property @@ -994,9 +1001,12 @@ class _ArrayOrScalarCommon: def __str__(self) -> str: ... def __repr__(self) -> str: ... def __copy__(self: _ArraySelf) -> _ArraySelf: ... - def __eq__(self, other): ... - def __ne__(self, other): ... def __deepcopy__(self: _ArraySelf, memo: None | Dict[int, Any], /) -> _ArraySelf: ... + + # TODO: How to deal with the non-commutative nature of `==` and `!=`? + # xref numpy/numpy#17368 + def __eq__(self, other: Any) -> Any: ... + def __ne__(self, other: Any) -> Any: ... def copy(self: _ArraySelf, order: _OrderKACF = ...) -> _ArraySelf: ... def dump(self, file: str | bytes | os.PathLike[str] | os.PathLike[bytes] | _SupportsWrite[bytes]) -> None: ... def dumps(self) -> bytes: ... @@ -1013,12 +1023,18 @@ class _ArrayOrScalarCommon: def tolist(self) -> Any: ... @property - def __array_interface__(self): ... + def __array_interface__(self) -> Dict[str, Any]: ... @property def __array_priority__(self) -> float: ... @property - def __array_struct__(self): ... - def __setstate__(self, state, /): ... + def __array_struct__(self) -> Any: ... # builtins.PyCapsule + def __setstate__(self, state: Tuple[ + SupportsIndex, # version + _ShapeLike, # Shape + _DType_co, # DType + bool, # F-continuous + bytes | List[Any], # Data + ], /) -> None: ... # a `bool_` is returned when `keepdims=True` and `self` is a 0d array @overload @@ -3428,8 +3444,6 @@ class format_parser: byteorder: None | _ByteOrder = ..., ) -> None: ... -# TODO: field-lookup returns either a `recarray` or a `ndarray` -# depending on the field dtype class recarray(ndarray[_ShapeType, _DType_co]): # NOTE: While not strictly mandatory, we're demanding here that arguments # for the `format_parser`- and `dtype`-based dtype constructors are @@ -3468,7 +3482,34 @@ class recarray(ndarray[_ShapeType, _DType_co]): def __array_finalize__(self, obj: object) -> None: ... def __getattribute__(self, attr: str) -> Any: ... def __setattr__(self, attr: str, val: ArrayLike) -> None: ... - def __getitem__(self, indx): ... # TODO + @overload + def __getitem__(self, indx: Union[ + SupportsIndex, + _ArrayLikeInt_co, + Tuple[SupportsIndex | _ArrayLikeInt_co, ...], + ]) -> Any: ... + @overload + def __getitem__(self: recarray[Any, dtype[void]], indx: Union[ + None, + slice, + ellipsis, + SupportsIndex, + _ArrayLikeInt_co, + Tuple[None | slice | ellipsis | _ArrayLikeInt_co | SupportsIndex, ...], + ]) -> recarray[Any, _DType_co]: ... + @overload + def __getitem__(self, indx: Union[ + None, + slice, + ellipsis, + SupportsIndex, + _ArrayLikeInt_co, + Tuple[None | slice | ellipsis | _ArrayLikeInt_co | SupportsIndex, ...], + ]) -> ndarray[Any, _DType_co]: ... + @overload + def __getitem__(self, indx: str) -> NDArray[Any]: ... + @overload + def __getitem__(self, indx: list[str]) -> recarray[_ShapeType, dtype[record]]: ... @overload def field(self, attr: int | str, val: None = ...) -> Any: ... @overload @@ -3638,7 +3679,6 @@ class memmap(ndarray[_ShapeType, _DType_co]): array: memmap[_ShapeType, _DType_co], context: None | Tuple[ufunc, Tuple[Any, ...], int] = ..., ) -> Any: ... - def __getitem__(self, index): ... # TODO def flush(self) -> None: ... class vectorize: @@ -3736,6 +3776,7 @@ class poly1d: ) -> poly1d: ... class matrix(ndarray[_ShapeType, _DType_co]): + __array_priority__: ClassVar[float] def __new__( subtype, data: ArrayLike, @@ -3743,7 +3784,27 @@ class matrix(ndarray[_ShapeType, _DType_co]): copy: bool = ..., ) -> matrix[Any, Any]: ... def __array_finalize__(self, obj: NDArray[Any]) -> None: ... - def __getitem__(self, index): ... # TODO + + @overload + def __getitem__(self, key: Union[ + SupportsIndex, + _ArrayLikeInt_co, + Tuple[SupportsIndex | _ArrayLikeInt_co, ...], + ]) -> Any: ... + @overload + def __getitem__(self, key: Union[ + None, + slice, + ellipsis, + SupportsIndex, + _ArrayLikeInt_co, + Tuple[None | slice | ellipsis | _ArrayLikeInt_co | SupportsIndex, ...], + ]) -> matrix[Any, _DType_co]: ... + @overload + def __getitem__(self: NDArray[void], key: str) -> matrix[Any, dtype[Any]]: ... + @overload + def __getitem__(self: NDArray[void], key: list[str]) -> matrix[_ShapeType, dtype[void]]: ... + def __mul__(self, other: ArrayLike) -> matrix[Any, Any]: ... def __rmul__(self, other: ArrayLike) -> matrix[Any, Any]: ... def __imul__(self, other: ArrayLike) -> matrix[_ShapeType, _DType_co]: ... |