diff options
| -rw-r--r-- | numpy/core/numerictypes.pyi | 97 |
1 files changed, 66 insertions, 31 deletions
diff --git a/numpy/core/numerictypes.pyi b/numpy/core/numerictypes.pyi index fd4aa3fda..e99e1c500 100644 --- a/numpy/core/numerictypes.pyi +++ b/numpy/core/numerictypes.pyi @@ -1,16 +1,15 @@ import sys +import types from typing import ( - TypeVar, - Optional, Type, Union, Tuple, - Sequence, overload, Any, TypeVar, Dict, List, + Iterable, ) from numpy import ( @@ -48,15 +47,22 @@ from numpy.core._type_aliases import ( sctypes as sctypes, ) -from numpy.typing import DTypeLike, ArrayLike +from numpy.typing import DTypeLike, ArrayLike, _SupportsDType if sys.version_info >= (3, 8): - from typing import Literal, Protocol, TypedDict + from typing import Literal as L, Protocol, TypedDict else: - from typing_extensions import Literal, Protocol, TypedDict + from typing_extensions import Literal as L, Protocol, TypedDict _T = TypeVar("_T") -_ScalarType = TypeVar("_ScalarType", bound=generic) +_SCT = TypeVar("_SCT", bound=generic) + +# A paramtrizable subset of `npt.DTypeLike` +_DTypeLike = Union[ + Type[_SCT], + dtype[_SCT], + _SupportsDType[dtype[_SCT]], +] class _CastFunc(Protocol): def __call__( @@ -64,42 +70,71 @@ class _CastFunc(Protocol): ) -> ndarray[Any, dtype[Any]]: ... class _TypeCodes(TypedDict): - Character: Literal['c'] - Integer: Literal['bhilqp'] - UnsignedInteger: Literal['BHILQP'] - Float: Literal['efdg'] - Complex: Literal['FDG'] - AllInteger: Literal['bBhHiIlLqQpP'] - AllFloat: Literal['efdgFDG'] - Datetime: Literal['Mm'] - All: Literal['?bhilqpBHILQPefdgFDGSUVOMm'] + Character: L['c'] + Integer: L['bhilqp'] + UnsignedInteger: L['BHILQP'] + Float: L['efdg'] + Complex: L['FDG'] + AllInteger: L['bBhHiIlLqQpP'] + AllFloat: L['efdgFDG'] + Datetime: L['Mm'] + All: L['?bhilqpBHILQPefdgFDGSUVOMm'] class _typedict(Dict[Type[generic], _T]): def __getitem__(self, key: DTypeLike) -> _T: ... +if sys.version_info >= (3, 10): + _TypeTuple = Union[ + Type[Any], + types.Union, + Tuple[Union[Type[Any], types.Union, Tuple[Any, ...]], ...], + ] +else: + _TypeTuple = Union[ + Type[Any], + Tuple[Union[Type[Any], Tuple[Any, ...]], ...], + ] + __all__: List[str] -# TODO: Clean up the annotations for the 7 functions below +@overload +def maximum_sctype(t: _DTypeLike[_SCT]) -> Type[_SCT]: ... +@overload +def maximum_sctype(t: DTypeLike) -> Type[Any]: ... + +@overload +def issctype(rep: dtype[Any] | Type[Any]) -> bool: ... +@overload +def issctype(rep: object) -> L[False]: ... -def maximum_sctype(t: DTypeLike) -> dtype: ... -def issctype(rep: object) -> bool: ... @overload -def obj2sctype(rep: object) -> Optional[generic]: ... +def obj2sctype(rep: _DTypeLike[_SCT], default: None = ...) -> None | Type[_SCT]: ... @overload -def obj2sctype(rep: object, default: None) -> Optional[generic]: ... +def obj2sctype(rep: _DTypeLike[_SCT], default: _T) -> _T | Type[_SCT]: ... @overload -def obj2sctype( - rep: object, default: Type[_T] -) -> Union[generic, Type[_T]]: ... -def issubclass_(arg1: object, arg2: Union[object, Tuple[object, ...]]) -> bool: ... -def issubsctype( - arg1: Union[ndarray, DTypeLike], arg2: Union[ndarray, DTypeLike] -) -> bool: ... +def obj2sctype(rep: DTypeLike, default: None = ...) -> None | Type[Any]: ... +@overload +def obj2sctype(rep: DTypeLike, default: _T) -> _T | Type[Any]: ... +@overload +def obj2sctype(rep: object, default: None = ...) -> None: ... +@overload +def obj2sctype(rep: object, default: _T) -> _T: ... + +@overload +def issubclass_(arg1: Type[Any], arg2: _TypeTuple) -> bool: ... +@overload +def issubclass_(arg1: object, arg2: object) -> L[False]: ... + +def issubsctype(arg1: DTypeLike, arg2: DTypeLike) -> bool: ... + def issubdtype(arg1: DTypeLike, arg2: DTypeLike) -> bool: ... -def sctype2char(sctype: object) -> str: ... + +def sctype2char(sctype: DTypeLike) -> str: ... + def find_common_type( - array_types: Sequence[DTypeLike], scalar_types: Sequence[DTypeLike] -) -> dtype: ... + array_types: Iterable[DTypeLike], + scalar_types: Iterable[DTypeLike], +) -> dtype[Any]: ... cast: _typedict[_CastFunc] nbytes: _typedict[int] |