From 553b9d15936704d9c96d272ee6ed5f13120ca310 Mon Sep 17 00:00:00 2001 From: jorenham Date: Mon, 24 Feb 2025 03:18:43 +0100 Subject: [PATCH] =?UTF-8?q?=F0=9F=90=9B=E2=9C=A8=20fix=20and=20improve=20`?= =?UTF-8?q?numpy.lib.=5Fpolynomial=5Fimpl`?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/numpy-stubs/__init__.pyi | 93 +--- src/numpy-stubs/lib/_polynomial_impl.pyi | 645 +++++++++++++++++++---- test/static/accept/lib_polynomial.pyi | 68 +-- 3 files changed, 567 insertions(+), 239 deletions(-) diff --git a/src/numpy-stubs/__init__.pyi b/src/numpy-stubs/__init__.pyi index 055aa8b9..2e6a5f48 100644 --- a/src/numpy-stubs/__init__.pyi +++ b/src/numpy-stubs/__init__.pyi @@ -239,7 +239,6 @@ from numpy._typing import ( _Float16Codes, _Float32Codes, _Float64Codes, - _FloatLike_co, _FloatingCodes, _GUFunc_Nin2_Nout1, _GenericCodes, @@ -407,7 +406,7 @@ from numpy.lib._nanfunctions_impl import ( nanvar, ) from numpy.lib._npyio_impl import fromregex, genfromtxt, load, loadtxt, save, savetxt, savez, savez_compressed -from numpy.lib._polynomial_impl import poly, polyadd, polyder, polydiv, polyfit, polyint, polymul, polysub, polyval, roots +from numpy.lib._polynomial_impl import poly, poly1d, polyadd, polyder, polydiv, polyfit, polyint, polymul, polysub, polyval, roots from numpy.lib._shape_base_impl import ( apply_along_axis, apply_over_axes, @@ -5261,96 +5260,6 @@ class vectorize: # def __call__(self, /, *args: Any, **kwargs: Any) -> Any: ... -class poly1d: - @property - def variable(self) -> LiteralString: ... - @property - def order(self) -> int: ... - @property - def o(self) -> int: ... - @property - def roots(self) -> NDArray[Any]: ... - @property - def r(self) -> NDArray[Any]: ... - - # - @property - def coeffs(self) -> NDArray[Any]: ... - @coeffs.setter - def coeffs(self, value: NDArray[Any]) -> None: ... - - # - @property - def c(self) -> NDArray[Any]: ... - @c.setter - def c(self, value: NDArray[Any]) -> None: ... - - # - @property - def coef(self) -> NDArray[Any]: ... - @coef.setter - def coef(self, value: NDArray[Any]) -> None: ... - - # - @property - def coefficients(self) -> NDArray[Any]: ... - @coefficients.setter - def coefficients(self, value: NDArray[Any]) -> None: ... - - __hash__: ClassVar[None] # type: ignore[assignment] # pyright: ignore[reportIncompatibleMethodOverride] - - # - def __init__(self, c_or_r: ArrayLike, r: builtins.bool = ..., variable: str | None = ...) -> None: ... - @overload - def __array__(self, /, t: None = None, copy: builtins.bool | None = None) -> ndarray[tuple[int], dtype[Any]]: ... - @overload - def __array__(self, /, t: _DType, copy: builtins.bool | None = None) -> ndarray[tuple[int], _DType]: ... - - # - @overload - def __call__(self, val: _ScalarLike_co) -> Any: ... - @overload - def __call__(self, val: poly1d) -> poly1d: ... - @overload - def __call__(self, val: ArrayLike) -> NDArray[Any]: ... - - # - def __len__(self) -> int: ... - def __iter__(self) -> Iterator[Any]: ... - - # - def __getitem__(self, val: int, /) -> Any: ... - def __setitem__(self, key: int, val: Any, /) -> None: ... - - # - def __neg__(self) -> Self: ... - def __pos__(self) -> Self: ... - - # - def __add__(self, other: ArrayLike, /) -> poly1d: ... - def __radd__(self, other: ArrayLike, /) -> poly1d: ... - - # - def __mul__(self, other: ArrayLike, /) -> poly1d: ... - def __rmul__(self, other: ArrayLike, /) -> poly1d: ... - - # - def __sub__(self, other: ArrayLike, /) -> poly1d: ... - def __rsub__(self, other: ArrayLike, /) -> poly1d: ... - - # - def __pow__(self, val: _FloatLike_co, /) -> poly1d: ... # Integral floats are accepted - - # - def __div__(self, other: ArrayLike, /) -> poly1d: ... - def __rdiv__(self, other: ArrayLike, /) -> poly1d: ... - def __truediv__(self, other: ArrayLike, /) -> poly1d: ... - def __rtruediv__(self, other: ArrayLike, /) -> poly1d: ... - - # - def deriv(self, m: SupportsInt | SupportsIndex = ...) -> poly1d: ... - def integ(self, m: SupportsInt | SupportsIndex = ..., k: _ArrayLikeComplex_co | _ArrayLikeObject_co | None = ...) -> Self: ... - class matrix(ndarray[_2DShapeT_co, _DType_co]): __array_priority__: ClassVar[float] = ... diff --git a/src/numpy-stubs/lib/_polynomial_impl.pyi b/src/numpy-stubs/lib/_polynomial_impl.pyi index 8ce980a4..9d391de0 100644 --- a/src/numpy-stubs/lib/_polynomial_impl.pyi +++ b/src/numpy-stubs/lib/_polynomial_impl.pyi @@ -1,178 +1,589 @@ -from typing import Any, Literal as L, NoReturn, SupportsIndex, SupportsInt, TypeAlias, overload -from typing_extensions import TypeVar +from collections.abc import Iterator +from typing import Any, ClassVar, Generic, Literal as L, SupportsIndex, SupportsInt, TypeAlias, overload +from typing_extensions import LiteralString, Self, TypeVar import numpy as np -from numpy import poly1d # noqa: ICN003 -from numpy._typing import ( - ArrayLike, - NDArray, - _ArrayLikeBool_co, - _ArrayLikeComplex_co, - _ArrayLikeFloat_co, - _ArrayLikeInt_co, - _ArrayLikeObject_co, - _ArrayLikeUInt_co, +from _numtype import ( + Array, + Array_1d, + CoComplex128_0d, + CoComplex128_1d, + CoComplex128_2d, + CoComplex_0d, + CoComplex_1d, + CoComplex_1nd, + CoComplex_2d, + CoComplex_nd, + CoFloat64_0d, + CoFloat64_1d, + CoFloating_0d, + CoFloating_1d, + CoFloating_1nd, + CoFloating_2d, + CoFloating_nd, + CoInteger_0d, + CoInteger_1d, + CoInteger_1nd, + CoInteger_2d, + CoSInteger_0d, + CoSInteger_1d, + CoSInteger_1nd, + CoSInteger_nd, + CoUInteger_0d, + CoUInteger_1d, + CoUInteger_1nd, + CoUInteger_nd, + Is, + ToBool_0d, + ToBool_1d, + ToBool_1nd, + ToBool_nd, + ToCLongDouble_1d, + ToComplex64_0d, + ToComplex64_1d, + ToComplex64_2d, + ToComplex128_0d, + ToComplex128_1d, + ToComplex128_2d, + ToComplex_0d, + ToComplex_1d, + ToComplex_1nd, + ToComplex_2d, + ToComplex_nd, + ToFloat32_1d, + ToFloat32_2d, + ToFloat64_1d, + ToFloat64_2d, + ToFloating_0d, + ToFloating_1d, + ToFloating_1nd, + ToFloating_nd, + ToInteger_0d, + ToInteger_1nd, + ToLongDouble_1d, + ToObject_1d, + ToObject_1nd, + ToObject_2d, + ToObject_nd, + ToSInteger_0d, + ToSInteger_1d, + ToSInteger_1nd, + ToSInteger_nd, + ToUInteger_0d, + ToUInteger_1d, + ToUInteger_1nd, + ToUInteger_nd, + _ToArray1_1d, + _ToArray_1d, ) +from numpy._typing import ArrayLike + +__all__ = ["poly", "poly1d", "polyder", "polydiv", "polyfit", "polyint", "polymul", "polymul", "polysub", "polyval", "roots"] + +### _T = TypeVar("_T") +_ScalarT = TypeVar("_ScalarT", bound=np.number | np.bool | np.object_) +_NumberT = TypeVar("_NumberT", bound=np.number) +_ScalarT_co = TypeVar("_ScalarT_co", bound=np.number | np.bool | np.object_, default=Any, covariant=True) + +_ToInt: TypeAlias = SupportsInt | SupportsIndex +_Tuple2: TypeAlias = tuple[_T, _T] +_Tuple_didd: TypeAlias = tuple[_T, Array[np.float64], Array[np.int32], Array[np.float64], Array[np.float64]] + +### + +class poly1d(Generic[_ScalarT_co]): + __hash__: ClassVar[None] # type: ignore[assignment] # pyright: ignore[reportIncompatibleMethodOverride] + + @property + def variable(self) -> LiteralString: ... + @property + def order(self) -> int: ... + @property + def o(self) -> int: ... + @property + def roots(self) -> Array_1d[np.inexact]: ... + @property + def r(self) -> Array_1d[np.inexact]: ... + + # + @property + def coefficients(self) -> Array_1d[_ScalarT_co]: ... + @coefficients.setter + def coefficients(self: poly1d[_ScalarT], c: Array[_ScalarT], /) -> None: ... + + # + @property + def coeffs(self) -> Array_1d[_ScalarT_co]: ... + @coeffs.setter + def coeffs(self: poly1d[_ScalarT], c: Array[_ScalarT], /) -> None: ... + + # + @property + def coef(self) -> Array_1d[_ScalarT_co]: ... + @coef.setter + def coef(self: poly1d[_ScalarT], c: Array[_ScalarT], /) -> None: ... + + # + @property + def c(self) -> Array_1d[_ScalarT_co]: ... + @c.setter + def c(self: poly1d[_ScalarT], c: Array[_ScalarT], /) -> None: ... + + # + @overload + def __init__(self, /, c_or_r: poly1d[_ScalarT_co], r: bool = False, variable: str | None = None) -> None: ... + @overload + def __init__(self, /, c_or_r: _ToArray1_1d[_ScalarT_co], r: L[False] = False, variable: str | None = None) -> None: ... + @overload + def __init__(self, /, c_or_r: CoComplex_1d, r: L[True], variable: str | None = None) -> None: ... + + # + @overload + def __array__(self, /, t: None = None, copy: bool | None = None) -> Array_1d[_ScalarT_co]: ... + @overload + def __array__(self, /, t: np.dtype[_ScalarT], copy: bool | None = None) -> Array_1d[_ScalarT]: ... + + # + @overload + def __call__(self, /, val: poly1d) -> poly1d: ... + @overload + def __call__(self: poly1d[np.object_], /, val: CoComplex_0d) -> Any: ... + @overload + def __call__(self: poly1d[np.bool], /, val: ToBool_0d) -> np.bool: ... # type: ignore[overload-overlap] + @overload + def __call__(self: poly1d[np.integer], /, val: CoInteger_0d) -> np.integer: ... # type: ignore[overload-overlap] + @overload + def __call__(self: poly1d[np.integer | np.bool], /, val: ToInteger_0d) -> np.integer: ... # type: ignore[overload-overlap] + @overload + def __call__(self: poly1d[np.floating], /, val: CoFloating_0d) -> np.floating: ... # type: ignore[overload-overlap] + @overload + def __call__(self: poly1d[np.floating | np.integer | np.bool], /, val: ToFloating_0d) -> np.floating: ... # type: ignore[overload-overlap] + @overload + def __call__(self: poly1d[np.complexfloating], /, val: CoComplex_0d) -> np.complexfloating: ... # type: ignore[overload-overlap] + @overload + def __call__(self: poly1d[np.number | np.bool], /, val: ToComplex_0d) -> np.complexfloating: ... # type: ignore[overload-overlap] + @overload + def __call__(self: poly1d[np.object_], /, val: CoComplex_1nd | ToObject_1nd) -> Array[Any]: ... + @overload + def __call__(self: poly1d[np.bool], /, val: ToBool_1nd) -> Array[np.bool]: ... + @overload + def __call__(self: poly1d[np.integer], /, val: CoInteger_1nd) -> Array[np.integer]: ... + @overload + def __call__(self: poly1d[np.integer | np.bool], /, val: ToInteger_1nd) -> Array[np.integer]: ... + @overload + def __call__(self: poly1d[np.floating], /, val: CoFloating_1nd) -> Array[np.floating]: ... + @overload + def __call__(self: poly1d[np.floating | np.integer | np.bool], /, val: ToFloating_1nd) -> Array[np.floating]: ... + @overload + def __call__(self: poly1d[np.complexfloating], /, val: CoComplex_1nd) -> Array[np.complexfloating]: ... + @overload + def __call__(self: poly1d[np.number | np.bool], /, val: ToComplex_1nd) -> Array[np.complexfloating]: ... -_2Tup: TypeAlias = tuple[_T, _T] -_5Tup: TypeAlias = tuple[_T, NDArray[np.float64], NDArray[np.int32], NDArray[np.float64], NDArray[np.float64]] + # + def __len__(self) -> int: ... -__all__ = ["poly", "poly1d", "polyadd", "polyder", "polydiv", "polyfit", "polyint", "polymul", "polysub", "polyval", "roots"] + # + @overload + def __iter__(self: poly1d[np.object_]) -> Iterator[Any]: ... + @overload + def __iter__(self) -> Iterator[_ScalarT_co]: ... -def poly(seq_of_zeros: ArrayLike) -> NDArray[np.floating]: ... + # + @overload + def __getitem__(self: poly1d[np.object_], val: int, /) -> Any: ... + @overload + def __getitem__(self, val: int, /) -> _ScalarT_co: ... -# Returns either a float or complex array depending on the input values. -# See `np.linalg.eigvals`. -def roots(p: ArrayLike) -> NDArray[np.inexact]: ... + # + def __setitem__(self, key: int, val: object, /) -> None: ... + + # + def __pos__(self) -> Self: ... + @overload + def __neg__(self: poly1d[_NumberT]) -> poly1d[_NumberT]: ... + @overload + def __neg__(self: poly1d[np.object_]) -> poly1d[np.object_]: ... + + # + def __add__(self, other: ArrayLike, /) -> poly1d: ... + def __radd__(self, other: ArrayLike, /) -> poly1d: ... + + # + def __mul__(self, other: ArrayLike, /) -> poly1d: ... + def __rmul__(self, other: ArrayLike, /) -> poly1d: ... + + # + def __sub__(self, other: ArrayLike, /) -> poly1d: ... + def __rsub__(self, other: ArrayLike, /) -> poly1d: ... + + # + def __pow__(self, val: CoFloating_0d, /) -> poly1d: ... # Integral floats are accepted + + # + def __div__(self, other: ArrayLike, /) -> poly1d: ... + def __rdiv__(self, other: ArrayLike, /) -> poly1d: ... + def __truediv__(self, other: ArrayLike, /) -> poly1d: ... + def __rtruediv__(self, other: ArrayLike, /) -> poly1d: ... + + # + @overload + def deriv(self: poly1d[_NumberT], /, m: _ToInt = 1) -> poly1d[_NumberT]: ... + @overload + def deriv(self: poly1d[np.bool], /, m: _ToInt = 1) -> poly1d[np.intp]: ... + @overload + def deriv(self: poly1d[np.object_], /, m: _ToInt = 1) -> poly1d[np.object_]: ... + + # + def integ(self, /, m: _ToInt = 1, k: CoComplex_0d | CoComplex_1d | ToObject_1d | None = None) -> poly1d: ... + +### + +# @overload -def polyint( - p: np.poly1d, - m: SupportsInt | SupportsIndex = ..., - k: _ArrayLikeComplex_co | _ArrayLikeObject_co | None = ..., -) -> np.poly1d: ... +def poly(seq_of_zeros: poly1d) -> Array_1d[np.float64]: ... +@overload +def poly(seq_of_zeros: CoInteger_1d | CoInteger_2d) -> Array_1d[np.float64]: ... +@overload +def poly(seq_of_zeros: ToFloat64_1d | ToFloat64_2d) -> Array_1d[np.float64]: ... # type: ignore[overload-overlap] +@overload +def poly(seq_of_zeros: ToComplex128_1d | ToComplex128_2d) -> Array_1d[np.complex128]: ... # type: ignore[overload-overlap] +@overload +def poly(seq_of_zeros: ToFloat32_1d | ToFloat32_2d) -> Array_1d[np.float32]: ... +@overload +def poly(seq_of_zeros: ToComplex64_1d | ToComplex64_2d) -> Array_1d[np.complex64]: ... +@overload +def poly(seq_of_zeros: ToObject_1d | ToObject_2d) -> Array_1d[np.object_]: ... +@overload +def poly(seq_of_zeros: CoComplex128_1d | CoComplex128_2d) -> Array_1d[np.inexact]: ... + +# returns either a float or complex array depending on the input values. See `np.linalg.eigvals`. +@overload +def roots(p: CoInteger_1d) -> Array_1d[np.float64 | np.complex128]: ... +@overload +def roots(p: ToFloat64_1d) -> Array_1d[np.float64 | np.complex128]: ... +@overload +def roots(p: ToFloat32_1d) -> Array_1d[np.float32 | np.complex64]: ... +@overload +def roots(p: ToComplex128_1d) -> Array_1d[np.complex128]: ... +@overload +def roots(p: ToComplex64_1d) -> Array_1d[np.complex64]: ... +@overload +def roots(p: CoComplex128_1d) -> Array_1d[np.inexact]: ... + +# +@overload +def polyder(p: poly1d, m: _ToInt = 1) -> poly1d: ... +@overload +def polyder(p: CoInteger_1d, m: _ToInt = 1) -> Array_1d[np.intp]: ... +@overload +def polyder(p: _ToArray_1d[np.float64 | np.float32 | np.float16, Is[float]], m: _ToInt = 1) -> Array_1d[np.float64]: ... # type: ignore[overload-overlap] +@overload +def polyder(p: _ToArray_1d[np.complex128 | np.complex64, Is[complex]], m: _ToInt = 1) -> Array_1d[np.complex128]: ... # type: ignore[overload-overlap] +@overload +def polyder(p: ToLongDouble_1d, m: _ToInt = 1) -> Array_1d[np.longdouble]: ... # type: ignore[overload-overlap] +@overload +def polyder(p: ToCLongDouble_1d, m: _ToInt = 1) -> Array_1d[np.clongdouble]: ... # type: ignore[overload-overlap] +@overload +def polyder(p: ToObject_1d, m: _ToInt = 1) -> Array_1d[np.object_]: ... +@overload +def polyder(p: CoComplex128_1d, m: _ToInt = 1) -> Array_1d[np.complex128 | np.float64 | np.intp]: ... + +# +@overload +def polyint(p: poly1d, m: _ToInt = 1, k: CoComplex_nd | ToObject_nd | None = None) -> poly1d: ... +@overload +def polyint( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] + p: CoFloat64_1d, + m: _ToInt = 1, + k: CoFloat64_0d | CoFloat64_1d | None = None, +) -> Array_1d[np.float64]: ... +@overload +def polyint( # type: ignore[overload-overlap] + p: ToLongDouble_1d, + m: _ToInt = 1, + k: CoFloating_0d | CoFloating_1d | None = None, +) -> Array_1d[np.longdouble]: ... +@overload +def polyint( # type: ignore[overload-overlap] + p: ToComplex128_1d | ToComplex64_1d, + m: _ToInt = 1, + k: CoComplex128_0d | CoComplex128_1d | None = None, +) -> Array_1d[np.complex128]: ... @overload def polyint( - p: _ArrayLikeFloat_co, - m: SupportsInt | SupportsIndex = ..., - k: _ArrayLikeFloat_co | None = ..., -) -> NDArray[np.floating]: ... + p: CoComplex128_1d, + m: _ToInt = 1, + k: ToComplex128_0d | ToComplex128_1d | ToComplex64_0d | ToComplex64_1d | None = None, +) -> Array_1d[np.complex128]: ... @overload def polyint( - p: _ArrayLikeComplex_co, - m: SupportsInt | SupportsIndex = ..., - k: _ArrayLikeComplex_co | None = ..., -) -> NDArray[np.complexfloating]: ... + p: ToCLongDouble_1d, + m: _ToInt = 1, + k: CoComplex_0d | CoComplex_1d | None = None, +) -> Array_1d[np.clongdouble]: ... @overload def polyint( - p: _ArrayLikeObject_co, - m: SupportsInt | SupportsIndex = ..., - k: _ArrayLikeObject_co | None = ..., -) -> NDArray[np.object_]: ... + p: ToObject_1d, + m: _ToInt = 1, + k: CoComplex_0d | CoComplex_1d | ToObject_1d | None = None, +) -> Array_1d[np.object_]: ... + +# +@overload +def polyfit( + x: CoFloating_1d, + y: CoFloating_1d | CoFloating_2d, + deg: _ToInt, + rcond: float | None = None, + full: L[False] = False, + w: CoFloating_1d | None = None, + cov: L[False] = False, +) -> Array[np.float64]: ... @overload -def polyder(p: np.poly1d, m: SupportsInt | SupportsIndex = ...) -> np.poly1d: ... +def polyfit( + x: CoFloating_1d, + y: CoFloating_1d | CoFloating_2d, + deg: _ToInt, + rcond: float | None = None, + full: L[False] = False, + w: CoFloating_1d | None = None, + *, + cov: L[True, "unscaled"], +) -> _Tuple2[Array[np.float64]]: ... +@overload +def polyfit( + x: CoFloating_1d, + y: CoFloating_1d | CoFloating_2d, + deg: _ToInt, + rcond: float | None, + full: L[True], + w: CoFloating_1d | None = None, + cov: bool | L["unscaled"] = False, +) -> _Tuple_didd[Array[np.float64]]: ... @overload -def polyder(p: _ArrayLikeFloat_co, m: SupportsInt | SupportsIndex = ...) -> NDArray[np.floating]: ... +def polyfit( + x: CoFloating_1d, + y: CoFloating_1d | CoFloating_2d, + deg: _ToInt, + rcond: float | None = None, + *, + full: L[True], + w: CoFloating_1d | None = None, + cov: bool | L["unscaled"] = False, +) -> _Tuple_didd[Array[np.float64]]: ... @overload -def polyder(p: _ArrayLikeComplex_co, m: SupportsInt | SupportsIndex = ...) -> NDArray[np.complexfloating]: ... +def polyfit( + x: CoComplex_1d, + y: ToComplex_1d | ToComplex_2d, + deg: _ToInt, + rcond: float | None = None, + full: L[False] = False, + w: CoFloating_1d | None = None, + cov: L[False] = False, +) -> Array[np.complex128]: ... @overload -def polyder(p: _ArrayLikeObject_co, m: SupportsInt | SupportsIndex = ...) -> NDArray[np.object_]: ... +def polyfit( + x: CoComplex_1d, + y: ToComplex_1d | ToComplex_2d, + deg: _ToInt, + rcond: float | None = None, + full: L[False] = False, + w: CoFloating_1d | None = None, + *, + cov: L[True, "unscaled"], +) -> _Tuple2[Array[np.complex128]]: ... @overload def polyfit( - x: _ArrayLikeFloat_co, - y: _ArrayLikeFloat_co, - deg: SupportsIndex | SupportsInt, - rcond: float | None = ..., - full: L[False] = ..., - w: _ArrayLikeFloat_co | None = ..., - cov: L[False] = ..., -) -> NDArray[np.float64]: ... + x: CoComplex_1d, + y: ToComplex_1d | ToComplex_2d, + deg: _ToInt, + rcond: float | None, + full: L[True], + w: CoFloating_1d | None = None, + cov: bool | L["unscaled"] = False, +) -> _Tuple_didd[Array[np.complex128]]: ... @overload def polyfit( - x: _ArrayLikeComplex_co, - y: _ArrayLikeComplex_co, - deg: SupportsIndex | SupportsInt, - rcond: float | None = ..., - full: L[False] = ..., - w: _ArrayLikeFloat_co | None = ..., - cov: L[False] = ..., -) -> NDArray[np.complex128]: ... + x: CoComplex_1d, + y: ToComplex_1d | ToComplex_2d, + deg: _ToInt, + rcond: float | None = None, + *, + full: L[True], + w: CoFloating_1d | None = None, + cov: bool | L["unscaled"] = False, +) -> _Tuple_didd[Array[np.complex128]]: ... @overload def polyfit( - x: _ArrayLikeFloat_co, - y: _ArrayLikeFloat_co, - deg: SupportsIndex | SupportsInt, - rcond: float | None = ..., - full: L[False] = ..., - w: _ArrayLikeFloat_co | None = ..., - cov: L[True, "unscaled"] = ..., -) -> _2Tup[NDArray[np.float64]]: ... + x: ToComplex_1d, + y: CoComplex_1d | CoComplex_2d, + deg: _ToInt, + rcond: float | None = None, + full: L[False] = False, + w: CoFloating_1d | None = None, + cov: L[False] = False, +) -> Array[np.complex128]: ... @overload def polyfit( - x: _ArrayLikeComplex_co, - y: _ArrayLikeComplex_co, - deg: SupportsIndex | SupportsInt, - rcond: float | None = ..., - full: L[False] = ..., - w: _ArrayLikeFloat_co | None = ..., - cov: L[True, "unscaled"] = ..., -) -> _2Tup[NDArray[np.complex128]]: ... + x: ToComplex_1d, + y: CoComplex_1d | CoComplex_2d, + deg: _ToInt, + rcond: float | None = None, + full: L[False] = False, + w: CoFloating_1d | None = None, + *, + cov: L[True, "unscaled"], +) -> _Tuple2[Array[np.complex128]]: ... @overload def polyfit( - x: _ArrayLikeFloat_co, - y: _ArrayLikeFloat_co, - deg: SupportsIndex | SupportsInt, - rcond: float | None = ..., - full: L[True] = ..., - w: _ArrayLikeFloat_co | None = ..., - cov: bool | L["unscaled"] = ..., -) -> _5Tup[NDArray[np.float64]]: ... + x: ToComplex_1d, + y: CoComplex_1d | CoComplex_2d, + deg: _ToInt, + rcond: float | None, + full: L[True], + w: CoFloating_1d | None = None, + cov: bool | L["unscaled"] = False, +) -> _Tuple_didd[Array[np.complex128]]: ... @overload def polyfit( - x: _ArrayLikeComplex_co, - y: _ArrayLikeComplex_co, - deg: SupportsIndex | SupportsInt, - rcond: float | None = ..., - full: L[True] = ..., - w: _ArrayLikeFloat_co | None = ..., - cov: bool | L["unscaled"] = ..., -) -> _5Tup[NDArray[np.complex128]]: ... + x: ToComplex_1d, + y: CoComplex_1d | CoComplex_2d, + deg: _ToInt, + rcond: float | None = None, + *, + full: L[True], + w: CoFloating_1d | None = None, + cov: bool | L["unscaled"] = False, +) -> _Tuple_didd[Array[np.complex128]]: ... + +# +@overload +def polyval(p: ToBool_1d, x: ToBool_0d) -> np.bool: ... # type: ignore[overload-overlap] +@overload +def polyval(p: ToUInteger_1d, x: CoUInteger_0d) -> np.unsignedinteger: ... # type: ignore[overload-overlap] +@overload +def polyval(p: CoUInteger_1d, x: ToUInteger_0d) -> np.unsignedinteger: ... # type: ignore[overload-overlap] +@overload +def polyval(p: ToSInteger_1d, x: CoSInteger_0d) -> np.signedinteger: ... # type: ignore[overload-overlap] +@overload +def polyval(p: CoSInteger_1d, x: ToSInteger_0d) -> np.signedinteger: ... # type: ignore[overload-overlap] +@overload +def polyval(p: ToFloating_1d, x: CoFloating_0d) -> np.floating: ... # type: ignore[overload-overlap] +@overload +def polyval(p: CoFloating_1d, x: ToFloating_0d) -> np.floating: ... # type: ignore[overload-overlap] +@overload +def polyval(p: ToComplex_1d, x: CoComplex_0d) -> np.complexfloating: ... # type: ignore[overload-overlap] +@overload +def polyval(p: CoComplex_1d, x: ToComplex_0d) -> np.complexfloating: ... # type: ignore[overload-overlap] +@overload +def polyval(p: ToBool_1d, x: ToBool_1nd) -> Array[np.bool]: ... +@overload +def polyval(p: ToUInteger_1d, x: CoUInteger_1nd) -> Array[np.unsignedinteger]: ... @overload -def polyval(p: _ArrayLikeBool_co, x: _ArrayLikeBool_co) -> NDArray[np.int64]: ... +def polyval(p: CoUInteger_1d, x: ToUInteger_1nd) -> Array[np.unsignedinteger]: ... @overload -def polyval(p: _ArrayLikeUInt_co, x: _ArrayLikeUInt_co) -> NDArray[np.unsignedinteger]: ... +def polyval(p: ToSInteger_1d, x: CoSInteger_1nd) -> Array[np.signedinteger]: ... @overload -def polyval(p: _ArrayLikeInt_co, x: _ArrayLikeInt_co) -> NDArray[np.signedinteger]: ... +def polyval(p: CoSInteger_1d, x: ToSInteger_1nd) -> Array[np.signedinteger]: ... @overload -def polyval(p: _ArrayLikeFloat_co, x: _ArrayLikeFloat_co) -> NDArray[np.floating]: ... +def polyval(p: ToFloating_1d, x: CoFloating_1nd) -> Array[np.floating]: ... @overload -def polyval(p: _ArrayLikeComplex_co, x: _ArrayLikeComplex_co) -> NDArray[np.complexfloating]: ... +def polyval(p: CoFloating_1d, x: ToFloating_1nd) -> Array[np.floating]: ... @overload -def polyval(p: _ArrayLikeObject_co, x: _ArrayLikeObject_co) -> NDArray[np.object_]: ... +def polyval(p: ToComplex_1d, x: CoComplex_1nd) -> Array[np.complexfloating]: ... @overload -def polyadd(a1: np.poly1d, a2: _ArrayLikeComplex_co | _ArrayLikeObject_co) -> np.poly1d: ... +def polyval(p: CoComplex_1d, x: ToComplex_1nd) -> Array[np.complexfloating]: ... @overload -def polyadd(a1: _ArrayLikeComplex_co | _ArrayLikeObject_co, a2: np.poly1d) -> np.poly1d: ... +def polyval(p: ToObject_1d, x: ToObject_1nd) -> Array[np.object_]: ... + +# +@overload +def polyadd(a1: poly1d, a2: CoComplex_nd | ToObject_nd) -> poly1d: ... +@overload +def polyadd(a1: CoComplex_nd | ToObject_nd, a2: poly1d) -> poly1d: ... +@overload +def polyadd(a1: ToBool_nd, a2: ToBool_nd) -> Array[np.bool]: ... +@overload +def polyadd(a1: ToUInteger_nd, a2: CoUInteger_nd) -> Array[np.unsignedinteger]: ... +@overload +def polyadd(a1: CoUInteger_nd, a2: ToUInteger_nd) -> Array[np.unsignedinteger]: ... @overload -def polyadd(a1: _ArrayLikeBool_co, a2: _ArrayLikeBool_co) -> NDArray[np.bool]: ... +def polyadd(a1: ToSInteger_nd, a2: CoSInteger_nd) -> Array[np.signedinteger]: ... @overload -def polyadd(a1: _ArrayLikeUInt_co, a2: _ArrayLikeUInt_co) -> NDArray[np.unsignedinteger]: ... +def polyadd(a1: CoSInteger_nd, a2: ToSInteger_nd) -> Array[np.signedinteger]: ... +@overload +def polyadd(a1: ToFloating_nd, a2: CoFloating_nd) -> Array[np.floating]: ... +@overload +def polyadd(a1: CoFloating_nd, a2: ToFloating_nd) -> Array[np.floating]: ... +@overload +def polyadd(a1: ToComplex_nd, a2: CoComplex_nd) -> Array[np.complexfloating]: ... +@overload +def polyadd(a1: CoComplex_nd, a2: ToComplex_nd) -> Array[np.complexfloating]: ... +@overload +def polyadd(a1: ToObject_nd, a2: ToObject_nd) -> Array[np.object_]: ... + +# keep in sync with polymul @overload -def polyadd(a1: _ArrayLikeInt_co, a2: _ArrayLikeInt_co) -> NDArray[np.signedinteger]: ... +def polymul(a1: poly1d, a2: CoComplex_nd | ToObject_nd) -> poly1d: ... @overload -def polyadd(a1: _ArrayLikeFloat_co, a2: _ArrayLikeFloat_co) -> NDArray[np.floating]: ... +def polymul(a1: CoComplex_nd | ToObject_nd, a2: poly1d) -> poly1d: ... @overload -def polyadd(a1: _ArrayLikeComplex_co, a2: _ArrayLikeComplex_co) -> NDArray[np.complexfloating]: ... +def polymul(a1: ToBool_nd, a2: ToBool_nd) -> Array[np.bool]: ... @overload -def polyadd(a1: _ArrayLikeObject_co, a2: _ArrayLikeObject_co) -> NDArray[np.object_]: ... +def polymul(a1: ToUInteger_nd, a2: CoUInteger_nd) -> Array[np.unsignedinteger]: ... @overload -def polysub(a1: np.poly1d, a2: _ArrayLikeComplex_co | _ArrayLikeObject_co) -> np.poly1d: ... +def polymul(a1: CoUInteger_nd, a2: ToUInteger_nd) -> Array[np.unsignedinteger]: ... @overload -def polysub(a1: _ArrayLikeComplex_co | _ArrayLikeObject_co, a2: np.poly1d) -> np.poly1d: ... +def polymul(a1: ToSInteger_nd, a2: CoSInteger_nd) -> Array[np.signedinteger]: ... @overload -def polysub(a1: _ArrayLikeBool_co, a2: _ArrayLikeBool_co) -> NoReturn: ... +def polymul(a1: CoSInteger_nd, a2: ToSInteger_nd) -> Array[np.signedinteger]: ... @overload -def polysub(a1: _ArrayLikeUInt_co, a2: _ArrayLikeUInt_co) -> NDArray[np.unsignedinteger]: ... +def polymul(a1: ToFloating_nd, a2: CoFloating_nd) -> Array[np.floating]: ... @overload -def polysub(a1: _ArrayLikeInt_co, a2: _ArrayLikeInt_co) -> NDArray[np.signedinteger]: ... +def polymul(a1: CoFloating_nd, a2: ToFloating_nd) -> Array[np.floating]: ... @overload -def polysub(a1: _ArrayLikeFloat_co, a2: _ArrayLikeFloat_co) -> NDArray[np.floating]: ... +def polymul(a1: ToComplex_nd, a2: CoComplex_nd) -> Array[np.complexfloating]: ... @overload -def polysub(a1: _ArrayLikeComplex_co, a2: _ArrayLikeComplex_co) -> NDArray[np.complexfloating]: ... +def polymul(a1: CoComplex_nd, a2: ToComplex_nd) -> Array[np.complexfloating]: ... @overload -def polysub(a1: _ArrayLikeObject_co, a2: _ArrayLikeObject_co) -> NDArray[np.object_]: ... +def polymul(a1: ToObject_nd, a2: ToObject_nd) -> Array[np.object_]: ... -# NOTE: Not an alias, but they do have the same signature (that we can reuse) -polymul = polyadd +# +@overload +def polysub(a1: poly1d, a2: CoComplex_nd | ToObject_nd) -> poly1d: ... +@overload +def polysub(a1: CoComplex_nd | ToObject_nd, a2: poly1d) -> poly1d: ... +@overload +def polysub(a1: ToUInteger_nd, a2: CoUInteger_nd) -> Array[np.unsignedinteger]: ... +@overload +def polysub(a1: CoUInteger_nd, a2: ToUInteger_nd) -> Array[np.unsignedinteger]: ... +@overload +def polysub(a1: ToSInteger_nd, a2: CoSInteger_nd) -> Array[np.signedinteger]: ... +@overload +def polysub(a1: CoSInteger_nd, a2: ToSInteger_nd) -> Array[np.signedinteger]: ... +@overload +def polysub(a1: ToFloating_nd, a2: CoFloating_nd) -> Array[np.floating]: ... +@overload +def polysub(a1: CoFloating_nd, a2: ToFloating_nd) -> Array[np.floating]: ... +@overload +def polysub(a1: ToComplex_nd, a2: CoComplex_nd) -> Array[np.complexfloating]: ... +@overload +def polysub(a1: CoComplex_nd, a2: ToComplex_nd) -> Array[np.complexfloating]: ... +@overload +def polysub(a1: ToObject_nd, a2: ToObject_nd) -> Array[np.object_]: ... +# +@overload +def polydiv(u: poly1d, v: CoComplex_nd | ToObject_nd) -> _Tuple2[poly1d]: ... @overload -def polydiv(u: np.poly1d, v: _ArrayLikeComplex_co | _ArrayLikeObject_co) -> _2Tup[np.poly1d]: ... +def polydiv(u: CoComplex_nd | ToObject_nd, v: poly1d) -> _Tuple2[poly1d]: ... @overload -def polydiv(u: _ArrayLikeComplex_co | _ArrayLikeObject_co, v: np.poly1d) -> _2Tup[np.poly1d]: ... +def polydiv(u: CoFloating_nd, v: CoFloating_nd) -> _Tuple2[Array[np.floating]]: ... @overload -def polydiv(u: _ArrayLikeFloat_co, v: _ArrayLikeFloat_co) -> _2Tup[NDArray[np.floating]]: ... +def polydiv(u: ToComplex_nd, v: CoComplex_nd) -> _Tuple2[Array[np.complexfloating]]: ... @overload -def polydiv(u: _ArrayLikeComplex_co, v: _ArrayLikeComplex_co) -> _2Tup[NDArray[np.complexfloating]]: ... +def polydiv(u: CoComplex_nd, v: ToComplex_nd) -> _Tuple2[Array[np.complexfloating]]: ... @overload -def polydiv(u: _ArrayLikeObject_co, v: _ArrayLikeObject_co) -> _2Tup[NDArray[Any]]: ... +def polydiv(u: ToObject_nd, v: ToObject_nd) -> _Tuple2[Array[Any]]: ... diff --git a/test/static/accept/lib_polynomial.pyi b/test/static/accept/lib_polynomial.pyi index b36ec4a3..41d2c5a2 100644 --- a/test/static/accept/lib_polynomial.pyi +++ b/test/static/accept/lib_polynomial.pyi @@ -1,10 +1,17 @@ from collections.abc import Iterator -from typing import Any, NoReturn -from typing_extensions import LiteralString, assert_type +from typing import Any, TypeAlias +from typing_extensions import LiteralString, TypeVar, assert_type import numpy as np import numpy.typing as npt +### + +_ScalarT = TypeVar("_ScalarT", bound=np.generic) +_Array1D: TypeAlias = np.ndarray[tuple[int], np.dtype[_ScalarT]] + +### + AR_b: npt.NDArray[np.bool] AR_u4: npt.NDArray[np.uint32] AR_i8: npt.NDArray[np.int64] @@ -14,20 +21,22 @@ AR_O: npt.NDArray[np.object_] poly_obj: np.poly1d +### + assert_type(poly_obj.variable, LiteralString) assert_type(poly_obj.order, int) assert_type(poly_obj.o, int) -assert_type(poly_obj.roots, npt.NDArray[Any]) -assert_type(poly_obj.r, npt.NDArray[Any]) -assert_type(poly_obj.coeffs, npt.NDArray[Any]) -assert_type(poly_obj.c, npt.NDArray[Any]) -assert_type(poly_obj.coef, npt.NDArray[Any]) -assert_type(poly_obj.coefficients, npt.NDArray[Any]) +assert_type(poly_obj.roots, _Array1D[np.inexact]) +assert_type(poly_obj.r, _Array1D[np.inexact]) +assert_type(poly_obj.coeffs, _Array1D[Any]) +assert_type(poly_obj.c, _Array1D[Any]) +assert_type(poly_obj.coef, _Array1D[Any]) +assert_type(poly_obj.coefficients, _Array1D[Any]) assert_type(poly_obj.__hash__, None) assert_type(poly_obj(1), Any) assert_type(poly_obj([1]), npt.NDArray[Any]) -assert_type(poly_obj(poly_obj), np.poly1d) +assert_type(poly_obj(poly_obj), np.poly1d) # type: ignore[assert-type] # mypy fail assert_type(len(poly_obj), int) assert_type(-poly_obj, np.poly1d) @@ -50,19 +59,19 @@ assert_type(iter(poly_obj), Iterator[Any]) assert_type(poly_obj.deriv(), np.poly1d) assert_type(poly_obj.integ(), np.poly1d) -assert_type(np.poly(poly_obj), npt.NDArray[np.floating]) -assert_type(np.poly(AR_f8), npt.NDArray[np.floating]) -assert_type(np.poly(AR_c16), npt.NDArray[np.floating]) +assert_type(np.poly(poly_obj), _Array1D[np.float64]) # type: ignore[assert-type] # mypy fail +assert_type(np.poly(AR_f8), _Array1D[np.float64]) +assert_type(np.poly(AR_c16), _Array1D[np.complex128]) -assert_type(np.polyint(poly_obj), np.poly1d) -assert_type(np.polyint(AR_f8), npt.NDArray[np.floating]) -assert_type(np.polyint(AR_f8, k=AR_c16), npt.NDArray[np.complexfloating]) -assert_type(np.polyint(AR_O, m=2), npt.NDArray[np.object_]) +assert_type(np.polyint(poly_obj), np.poly1d) # type: ignore[assert-type] # mypy fail +assert_type(np.polyint(AR_f8), _Array1D[np.float64]) +assert_type(np.polyint(AR_f8, k=AR_c16), _Array1D[np.complex128]) +assert_type(np.polyint(AR_O, m=2), _Array1D[np.object_]) -assert_type(np.polyder(poly_obj), np.poly1d) -assert_type(np.polyder(AR_f8), npt.NDArray[np.floating]) -assert_type(np.polyder(AR_c16), npt.NDArray[np.complexfloating]) -assert_type(np.polyder(AR_O, m=2), npt.NDArray[np.object_]) +assert_type(np.polyder(poly_obj), np.poly1d) # type: ignore[assert-type] # mypy fail +assert_type(np.polyder(AR_f8), _Array1D[np.float64]) +assert_type(np.polyder(AR_c16), _Array1D[np.complex128]) +assert_type(np.polyder(AR_O, m=2), _Array1D[np.object_]) assert_type(np.polyfit(AR_f8, AR_f8, 2), npt.NDArray[np.float64]) assert_type( @@ -101,15 +110,15 @@ assert_type( ], ) -assert_type(np.polyval(AR_b, AR_b), npt.NDArray[np.int64]) +assert_type(np.polyval(AR_b, AR_b), npt.NDArray[np.bool]) assert_type(np.polyval(AR_u4, AR_b), npt.NDArray[np.unsignedinteger]) assert_type(np.polyval(AR_i8, AR_i8), npt.NDArray[np.signedinteger]) assert_type(np.polyval(AR_f8, AR_i8), npt.NDArray[np.floating]) assert_type(np.polyval(AR_i8, AR_c16), npt.NDArray[np.complexfloating]) assert_type(np.polyval(AR_O, AR_O), npt.NDArray[np.object_]) -assert_type(np.polyadd(poly_obj, AR_i8), np.poly1d) -assert_type(np.polyadd(AR_f8, poly_obj), np.poly1d) +assert_type(np.polyadd(poly_obj, AR_i8), np.poly1d) # type: ignore[assert-type] # mypy fail +assert_type(np.polyadd(AR_f8, poly_obj), np.poly1d) # type: ignore[assert-type] # mypy fail assert_type(np.polyadd(AR_b, AR_b), npt.NDArray[np.bool]) assert_type(np.polyadd(AR_u4, AR_b), npt.NDArray[np.unsignedinteger]) assert_type(np.polyadd(AR_i8, AR_i8), npt.NDArray[np.signedinteger]) @@ -117,17 +126,16 @@ assert_type(np.polyadd(AR_f8, AR_i8), npt.NDArray[np.floating]) assert_type(np.polyadd(AR_i8, AR_c16), npt.NDArray[np.complexfloating]) assert_type(np.polyadd(AR_O, AR_O), npt.NDArray[np.object_]) -assert_type(np.polysub(poly_obj, AR_i8), np.poly1d) -assert_type(np.polysub(AR_f8, poly_obj), np.poly1d) -assert_type(np.polysub(AR_b, AR_b), NoReturn) +assert_type(np.polysub(poly_obj, AR_i8), np.poly1d) # type: ignore[assert-type] # mypy fail +assert_type(np.polysub(AR_f8, poly_obj), np.poly1d) # type: ignore[assert-type] # mypy fail assert_type(np.polysub(AR_u4, AR_b), npt.NDArray[np.unsignedinteger]) assert_type(np.polysub(AR_i8, AR_i8), npt.NDArray[np.signedinteger]) assert_type(np.polysub(AR_f8, AR_i8), npt.NDArray[np.floating]) assert_type(np.polysub(AR_i8, AR_c16), npt.NDArray[np.complexfloating]) assert_type(np.polysub(AR_O, AR_O), npt.NDArray[np.object_]) -assert_type(np.polymul(poly_obj, AR_i8), np.poly1d) -assert_type(np.polymul(AR_f8, poly_obj), np.poly1d) +assert_type(np.polymul(poly_obj, AR_i8), np.poly1d) # type: ignore[assert-type] # mypy fail +assert_type(np.polymul(AR_f8, poly_obj), np.poly1d) # type: ignore[assert-type] # mypy fail assert_type(np.polymul(AR_b, AR_b), npt.NDArray[np.bool]) assert_type(np.polymul(AR_u4, AR_b), npt.NDArray[np.unsignedinteger]) assert_type(np.polymul(AR_i8, AR_i8), npt.NDArray[np.signedinteger]) @@ -135,8 +143,8 @@ assert_type(np.polymul(AR_f8, AR_i8), npt.NDArray[np.floating]) assert_type(np.polymul(AR_i8, AR_c16), npt.NDArray[np.complexfloating]) assert_type(np.polymul(AR_O, AR_O), npt.NDArray[np.object_]) -assert_type(np.polydiv(poly_obj, AR_i8), tuple[np.poly1d, np.poly1d]) -assert_type(np.polydiv(AR_f8, poly_obj), tuple[np.poly1d, np.poly1d]) +assert_type(np.polydiv(poly_obj, AR_i8), tuple[np.poly1d, np.poly1d]) # type: ignore[assert-type] # mypy fail +assert_type(np.polydiv(AR_f8, poly_obj), tuple[np.poly1d, np.poly1d]) # type: ignore[assert-type] # mypy fail assert_type(np.polydiv(AR_b, AR_b), tuple[npt.NDArray[np.floating], npt.NDArray[np.floating]]) assert_type(np.polydiv(AR_u4, AR_b), tuple[npt.NDArray[np.floating], npt.NDArray[np.floating]]) assert_type(np.polydiv(AR_i8, AR_i8), tuple[npt.NDArray[np.floating], npt.NDArray[np.floating]])