Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions src/numpy-stubs/@test/static/accept/type_check.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -26,12 +26,12 @@ assert_type(np.mintypecode(["f8"], typeset="qfQF"), str)
assert_type(np.real(ComplexObj()), slice)
assert_type(np.real(AR_f8), _nt.Array[np.float64])
assert_type(np.real(AR_c16), _nt.Array[np.float64])
assert_type(np.real(AR_LIKE_i), _nt.Array[np.intp])
assert_type(np.real(AR_LIKE_i), _nt.Array[Any])

assert_type(np.imag(ComplexObj()), slice)
assert_type(np.imag(AR_f8), _nt.Array[np.float64])
assert_type(np.imag(AR_c16), _nt.Array[np.float64])
assert_type(np.imag(AR_LIKE_i), _nt.Array[np.intp])
assert_type(np.imag(AR_LIKE_i), _nt.Array[Any])

assert_type(np.iscomplex(f8), np.bool)
assert_type(np.iscomplex(AR_f8), _nt.Array[np.bool])
Expand Down
175 changes: 69 additions & 106 deletions src/numpy-stubs/lib/_type_check_impl.pyi
Original file line number Diff line number Diff line change
@@ -1,9 +1,11 @@
from _typeshed import Incomplete
from collections.abc import Container, Iterable
from typing import Any, Literal as L, overload, type_check_only
from typing_extensions import Protocol, TypeVar
from typing import Any, Literal as L, Protocol, TypeAlias, overload, type_check_only
from typing_extensions import TypeVar

import _numtype as _nt
import numpy as np
from numpy._typing import ArrayLike, _ArrayLike, _NestedSequence, _ScalarLike_co, _SupportsArray

__all__ = [
"common_type",
Expand All @@ -19,14 +21,18 @@ __all__ = [
"typename",
]

###

_T = TypeVar("_T")
_T_co = TypeVar("_T_co", covariant=True)

_RealT = TypeVar("_RealT", bound=_nt.co_number)
_ScalarT = TypeVar("_ScalarT", bound=np.generic)
_ScalarT_co = TypeVar("_ScalarT_co", bound=np.generic, covariant=True)
_RealT = TypeVar("_RealT", bound=np.floating | np.integer | np.bool)

_FloatMax32: TypeAlias = np.float32 | np.float16
_ComplexMax128: TypeAlias = np.complex128 | np.complex64
_RealMax64: TypeAlias = np.float64 | np.float32 | np.float16 | np.integer
_Real: TypeAlias = np.floating | np.integer
_InexactMax32: TypeAlias = np.float32 | np.complex64 | np.float16
_NumberMax64: TypeAlias = np.float64 | np.complex128 | _InexactMax32 | np.integer

@type_check_only
class _HasReal(Protocol[_T_co]):
Expand All @@ -45,66 +51,45 @@ class _HasDType(Protocol[_ScalarT_co]):

###

#
def mintypecode(
typechars: Iterable[_nt.ToGeneric_nd], typeset: str | Container[str] = "GDFgdf", default: str = "d"
typechars: Iterable[str | ArrayLike], typeset: str | Container[str] = "GDFgdf", default: str = "d"
) -> str: ...

#
@overload
def real(val: _HasReal[_T]) -> _T: ... # type: ignore[overload-overlap]
@overload
def real(val: _nt.ToBool_nd) -> _nt.Array[np.bool]: ...
@overload
def real(val: _nt.ToInt_nd) -> _nt.Array[np.intp]: ...
@overload
def real(val: _nt.ToFloat64_nd) -> _nt.Array[np.float64]: ...
@overload
def real(val: _nt.ToComplex128_nd) -> _nt.Array[np.complex128]: ...
def real(val: _HasReal[_T]) -> _T: ... # type: ignore[overload-overlap] # false positive
@overload
def real(val: _nt.ToBytes_nd) -> _nt.Array[np.bytes_]: ...
def real(val: _ArrayLike[_RealT]) -> _nt.Array[_RealT]: ...
@overload
def real(val: _nt.ToStr_nd) -> _nt.Array[np.str_]: ...
@overload
def real(val: _nt._ToArray_nd[_ScalarT]) -> _nt.Array[_ScalarT]: ...
@overload
def real(val: _nt.ToGeneric_nd) -> _nt.Array: ...
def real(val: ArrayLike) -> _nt.Array[Any]: ...

#
@overload
def imag(val: _HasImag[_T]) -> _T: ... # type: ignore[overload-overlap]
@overload
def imag(val: _nt.ToBool_nd) -> _nt.Array[np.bool]: ...
@overload
def imag(val: _nt.ToInt_nd) -> _nt.Array[np.intp]: ...
@overload
def imag(val: _nt.ToFloat64_nd) -> _nt.Array[np.float64]: ...
def imag(val: _HasImag[_T]) -> _T: ... # type: ignore[overload-overlap] # false positive
@overload
def imag(val: _nt.ToComplex128_nd) -> _nt.Array[np.complex128]: ...
def imag(val: _ArrayLike[_RealT]) -> _nt.Array[_RealT]: ...
@overload
def imag(val: _nt.ToBytes_nd) -> _nt.Array[np.bytes_]: ...
@overload
def imag(val: _nt.ToStr_nd) -> _nt.Array[np.str_]: ...
@overload
def imag(val: _nt._ToArray_nd[_ScalarT]) -> _nt.Array[_ScalarT]: ...
@overload
def imag(val: _nt.ToGeneric_nd) -> _nt.Array: ...
def imag(val: ArrayLike) -> _nt.Array[Any]: ...

#
@overload
def iscomplex(x: _nt.ToGeneric_1nd) -> _nt.Array[np.bool]: ...
def iscomplex(x: _ScalarLike_co) -> np.bool: ...
@overload
def iscomplex(x: np.ndarray | _NestedSequence[ArrayLike]) -> _nt.Array[np.bool]: ...
@overload
def iscomplex(x: np.generic | _nt._PyScalar) -> np.bool: ...
def iscomplex(x: ArrayLike) -> np.bool | _nt.Array[np.bool]: ...

#
@overload
def isreal(x: _nt.ToGeneric_1nd) -> _nt.Array[np.bool]: ...
def isreal(x: _ScalarLike_co) -> np.bool: ...
@overload
def isreal(x: np.generic | _nt._PyScalar) -> np.bool: ...
def isreal(x: np.ndarray | _NestedSequence[ArrayLike]) -> _nt.Array[np.bool]: ...
@overload
def isreal(x: ArrayLike) -> np.bool | _nt.Array[np.bool]: ...

#
def iscomplexobj(x: _HasDType[Any] | _nt.ToGeneric_nd) -> bool: ...
def isrealobj(x: _HasDType[Any] | _nt.ToGeneric_nd) -> bool: ...
def iscomplexobj(x: _HasDType[Any] | ArrayLike) -> bool: ...
def isrealobj(x: _HasDType[Any] | ArrayLike) -> bool: ...

#
@overload
Expand All @@ -113,34 +98,44 @@ def nan_to_num(
) -> _ScalarT: ...
@overload
def nan_to_num(
x: _nt._ToArray_1nd[_ScalarT],
x: _nt.Array[_ScalarT] | _NestedSequence[_ArrayLike[_ScalarT]],
copy: bool = True,
nan: float = 0.0,
posinf: float | None = None,
neginf: float | None = None,
) -> _nt.Array[_ScalarT]: ...
@overload
def nan_to_num(
x: _nt.ToGeneric_0d, copy: bool = True, nan: float = 0.0, posinf: float | None = None, neginf: float | None = None
) -> Any: ...
x: _SupportsArray[np.dtype[_ScalarT]],
copy: bool = True,
nan: float = 0.0,
posinf: float | None = None,
neginf: float | None = None,
) -> _ScalarT | _nt.Array[_ScalarT]: ...
@overload
def nan_to_num(
x: _nt.ToGeneric_1nd, copy: bool = True, nan: float = 0.0, posinf: float | None = None, neginf: float | None = None
) -> _nt.Array: ...

# If one passes a complex array to `real_if_close`, then one is reasonably
# expected to verify the output dtype (so we can return an unsafe union here)
x: _NestedSequence[ArrayLike],
copy: bool = True,
nan: float = 0.0,
posinf: float | None = None,
neginf: float | None = None,
) -> _nt.Array[Incomplete]: ...
@overload
def nan_to_num(
x: ArrayLike, copy: bool = True, nan: float = 0.0, posinf: float | None = None, neginf: float | None = None
) -> Incomplete: ...

# NOTE: The [overload-overlap] mypy error is a false positive
@overload
def real_if_close(a: _nt.ToComplex128_nd, tol: float = 100) -> _nt.Array[_nt.inexact64]: ...
def real_if_close(a: _ArrayLike[np.complex64], tol: float = 100) -> _nt.Array[np.float32 | np.complex64]: ...
@overload
def real_if_close(a: _nt.ToComplex64_nd, tol: float = 100) -> _nt.Array[_nt.inexact32]: ...
def real_if_close(a: _ArrayLike[np.complex128], tol: float = 100) -> _nt.Array[np.float64 | np.complex128]: ...
@overload
def real_if_close(a: _nt.ToCLongDouble_nd, tol: float = 100) -> _nt.Array[_nt.inexact64l]: ...
def real_if_close(a: _ArrayLike[np.clongdouble], tol: float = 100) -> _nt.Array[np.longdouble | np.clongdouble]: ...
@overload
def real_if_close(a: _nt._ToArray_nd[_RealT], tol: float = 100) -> _nt.Array[_RealT]: ...
def real_if_close(a: _ArrayLike[_RealT], tol: float = 100) -> _nt.Array[_RealT]: ...
@overload
def real_if_close(a: _nt.ToGeneric_nd, tol: float = 100) -> _nt.Array: ...
def real_if_close(a: ArrayLike, tol: float = 100) -> _nt.Array[Any]: ...

#
@overload
Expand Down Expand Up @@ -188,91 +183,61 @@ def typename(char: L["V"]) -> L["void"]: ...
@overload
def typename(char: L["O"]) -> L["object"]: ...

# NOTE: mypy reports false-positive overlapping overloads
# NOTE: The [overload-overlap] mypy errors are false positives
@overload
def common_type() -> type[np.float16]: ...
@overload
def common_type(a0: _HasDType[np.float16], /, *ai: _HasDType[np.float16]) -> type[np.float16]: ...
@overload
def common_type(a0: _HasDType[np.float32], /, *ai: _HasDType[np.float32 | np.float16]) -> type[np.float32]: ...
def common_type(a0: _HasDType[np.float32], /, *ai: _HasDType[_FloatMax32]) -> type[np.float32]: ...
@overload
def common_type(
a0: _HasDType[np.float64 | np.integer], /, *ai: _HasDType[np.float64 | np.float32 | np.float16 | np.integer]
) -> type[np.float64]: ...
def common_type(a0: _HasDType[np.float64 | np.integer], /, *ai: _HasDType[_RealMax64]) -> type[np.float64]: ...
@overload
def common_type(a0: _HasDType[np.longdouble], /, *ai: _HasDType[np.floating | np.integer]) -> type[np.longdouble]: ...
def common_type(a0: _HasDType[np.longdouble], /, *ai: _HasDType[_Real]) -> type[np.longdouble]: ...
@overload
def common_type(a0: _HasDType[np.complex64], /, *ai: _HasDType[_nt.inexact32 | np.float16]) -> type[np.complex64]: ...
def common_type(a0: _HasDType[np.complex64], /, *ai: _HasDType[_InexactMax32]) -> type[np.complex64]: ...
@overload
def common_type(
a0: _HasDType[np.complex128], /, *ai: _HasDType[_nt.inexact64 | _nt.inexact32 | np.float16 | np.integer]
) -> type[np.complex128]: ...
def common_type(a0: _HasDType[np.complex128], /, *ai: _HasDType[_NumberMax64]) -> type[np.complex128]: ...
@overload
def common_type(a0: _HasDType[np.clongdouble], /, *ai: _HasDType[np.number]) -> type[np.clongdouble]: ...
@overload
def common_type(
a0: _HasDType[np.float32 | np.float16], array1: _HasDType[np.float32], /, *ai: _HasDType[np.float32 | np.float16]
a0: _HasDType[_FloatMax32], array1: _HasDType[np.float32], /, *ai: _HasDType[_FloatMax32]
) -> type[np.float32]: ...
@overload
def common_type(
a0: _HasDType[np.float64 | np.float32 | np.float16 | np.integer],
array1: _HasDType[np.float64 | np.integer],
/,
*ai: _HasDType[np.float64 | np.float32 | np.float16 | np.integer],
a0: _HasDType[_RealMax64], array1: _HasDType[np.float64 | np.integer], /, *ai: _HasDType[_RealMax64]
) -> type[np.float64]: ...
@overload
def common_type(
a0: _HasDType[np.floating | np.integer],
array1: _HasDType[np.longdouble],
/,
*ai: _HasDType[np.floating | np.integer],
a0: _HasDType[_Real], array1: _HasDType[np.longdouble], /, *ai: _HasDType[_Real]
) -> type[np.longdouble]: ...
@overload
def common_type(
a0: _HasDType[_nt.inexact32 | np.float16],
array1: _HasDType[np.complex64],
/,
*ai: _HasDType[_nt.inexact32 | np.float16],
a0: _HasDType[_InexactMax32], array1: _HasDType[np.complex64], /, *ai: _HasDType[_InexactMax32]
) -> type[np.complex64]: ...
@overload
def common_type(
a0: _HasDType[np.float64],
array1: _HasDType[np.complex128 | np.complex64],
/,
*ai: _HasDType[_nt.inexact64 | _nt.inexact32 | np.float16 | np.integer],
a0: _HasDType[np.float64], array1: _HasDType[_ComplexMax128], /, *ai: _HasDType[_NumberMax64]
) -> type[np.complex128]: ...
@overload
def common_type(
a0: _HasDType[np.complex128 | np.complex64],
array1: _HasDType[np.float64],
/,
*ai: _HasDType[_nt.inexact64 | _nt.inexact32 | np.float16 | np.integer],
a0: _HasDType[_ComplexMax128], array1: _HasDType[np.float64], /, *ai: _HasDType[_NumberMax64]
) -> type[np.complex128]: ...
@overload
def common_type(
a0: _HasDType[_nt.inexact64 | _nt.inexact32 | np.float16 | np.integer],
array1: _HasDType[np.complex128],
/,
*ai: _HasDType[_nt.inexact64 | _nt.inexact32 | np.float16 | np.integer],
a0: _HasDType[_NumberMax64], array1: _HasDType[np.complex128], /, *ai: _HasDType[_NumberMax64]
) -> type[np.complex128]: ...
@overload
def common_type(
a0: _HasDType[np.complex128 | np.complex64],
array1: _HasDType[np.complex128 | np.integer],
/,
*ai: _HasDType[_nt.inexact64 | _nt.inexact32 | np.float16 | np.integer],
a0: _HasDType[_ComplexMax128], array1: _HasDType[np.complex128 | np.integer], /, *ai: _HasDType[_NumberMax64]
) -> type[np.complex128]: ...
@overload
def common_type(
a0: _HasDType[np.complex128 | np.integer],
array1: _HasDType[np.complex128 | np.complex64],
/,
*ai: _HasDType[_nt.inexact64 | _nt.inexact32 | np.float16 | np.integer],
a0: _HasDType[np.complex128 | np.integer], array1: _HasDType[_ComplexMax128], /, *ai: _HasDType[_NumberMax64]
) -> type[np.complex128]: ...
@overload
def common_type(
a0: _HasDType[np.floating | np.integer], /, *ai: _HasDType[np.floating | np.integer]
) -> type[np.floating]: ...
def common_type(a0: _HasDType[_Real], /, *ai: _HasDType[_Real]) -> type[np.floating]: ...
@overload
def common_type(
a0: _HasDType[np.number], array1: _HasDType[np.clongdouble], /, *ai: _HasDType[np.number]
Expand All @@ -294,6 +259,4 @@ def common_type(
a0: _HasDType[np.number], array1: _HasDType[np.complexfloating], /, *ai: _HasDType[np.number]
) -> type[np.complexfloating]: ...
@overload
def common_type(
a0: _HasDType[np.number], array1: _HasDType[np.number], /, *ai: _HasDType[np.number]
) -> type[np.inexact]: ...
def common_type(a0: _HasDType[np.number], array1: _HasDType[np.number], /, *ai: _HasDType[np.number]) -> type[Any]: ...