ENH: Add improved placeholder annotations for `np.polynomial`

9 files changed, 370 insertions, 8 deletions

diff --git a/numpy/polynomial/__init__.pyi b/numpy/polynomial/__init__.pyi
index 6a7406041..bebedb3a6 100644
--- a/numpy/polynomial/__init__.pyi
+++ b/numpy/polynomial/__init__.pyi
@@ -1,4 +1,4 @@
-from typing import Any, List
+from typing import List
 
 from numpy.polynomial import (
     chebyshev as chebyshev,
@@ -8,13 +8,13 @@ from numpy.polynomial import (
     legendre as legendre,
     polynomial as polynomial,
 )
+from numpy.polynomial.chebyshev import Chebyshev as Chebyshev
+from numpy.polynomial.hermite import Hermite as Hermite
+from numpy.polynomial.hermite_e import HermiteE as HermiteE
+from numpy.polynomial.laguerre import Laguerre as Laguerre
+from numpy.polynomial.legendre import Legendre as Legendre
+from numpy.polynomial.polynomial import Polynomial as Polynomial
 
 __all__: List[str]
 
-Polynomial: Any
-Chebyshev: Any
-Legendre: Any
-Hermite: Any
-HermiteE: Any
-Laguerre: Any
-set_default_printstyle: Any
+def set_default_printstyle(style): ...
diff --git a/numpy/polynomial/_polybase.pyi b/numpy/polynomial/_polybase.pyi
new file mode 100644
index 000000000..c41601469
--- /dev/null
+++ b/numpy/polynomial/_polybase.pyi
@@ -0,0 +1,69 @@
+import abc
+from typing import Any, List, ClassVar
+
+__all__: List[str]
+
+class ABCPolyBase(abc.ABC):
+    __hash__: ClassVar[None]  # type: ignore[assignment]
+    __array_ufunc__: ClassVar[None]
+    maxpower: ClassVar[int]
+    coef: Any
+    @property
+    @abc.abstractmethod
+    def domain(self): ...
+    @property
+    @abc.abstractmethod
+    def window(self): ...
+    @property
+    @abc.abstractmethod
+    def basis_name(self): ...
+    def has_samecoef(self, other): ...
+    def has_samedomain(self, other): ...
+    def has_samewindow(self, other): ...
+    def has_sametype(self, other): ...
+    def __init__(self, coef, domain=..., window=...): ...
+    def __format__(self, fmt_str): ...
+    def __call__(self, arg): ...
+    def __iter__(self): ...
+    def __len__(self): ...
+    def __neg__(self): ...
+    def __pos__(self): ...
+    def __add__(self, other): ...
+    def __sub__(self, other): ...
+    def __mul__(self, other): ...
+    def __truediv__(self, other): ...
+    def __floordiv__(self, other): ...
+    def __mod__(self, other): ...
+    def __divmod__(self, other): ...
+    def __pow__(self, other): ...
+    def __radd__(self, other): ...
+    def __rsub__(self, other): ...
+    def __rmul__(self, other): ...
+    def __rdiv__(self, other): ...
+    def __rtruediv__(self, other): ...
+    def __rfloordiv__(self, other): ...
+    def __rmod__(self, other): ...
+    def __rdivmod__(self, other): ...
+    def __eq__(self, other): ...
+    def __ne__(self, other): ...
+    def copy(self): ...
+    def degree(self): ...
+    def cutdeg(self, deg): ...
+    def trim(self, tol=...): ...
+    def truncate(self, size): ...
+    def convert(self, domain=..., kind=..., window=...): ...
+    def mapparms(self): ...
+    def integ(self, m=..., k = ..., lbnd=...): ...
+    def deriv(self, m=...): ...
+    def roots(self): ...
+    def linspace(self, n=..., domain=...): ...
+    @classmethod
+    def fit(cls, x, y, deg, domain=..., rcond=..., full=..., w=..., window=...): ...
+    @classmethod
+    def fromroots(cls, roots, domain = ..., window=...): ...
+    @classmethod
+    def identity(cls, domain=..., window=...): ...
+    @classmethod
+    def basis(cls, deg, domain=..., window=...): ...
+    @classmethod
+    def cast(cls, series, domain=..., window=...): ...
diff --git a/numpy/polynomial/chebyshev.pyi b/numpy/polynomial/chebyshev.pyi
new file mode 100644
index 000000000..841c0859b
--- /dev/null
+++ b/numpy/polynomial/chebyshev.pyi
@@ -0,0 +1,51 @@
+from typing import Any, List
+
+from numpy import ndarray, dtype, int_
+from numpy.polynomial._polybase import ABCPolyBase
+from numpy.polynomial.polyutils import trimcoef
+
+__all__: List[str]
+
+chebtrim = trimcoef
+
+def poly2cheb(pol): ...
+def cheb2poly(c): ...
+
+chebdomain: ndarray[Any, dtype[int_]]
+chebzero: ndarray[Any, dtype[int_]]
+chebone: ndarray[Any, dtype[int_]]
+chebx: ndarray[Any, dtype[int_]]
+
+def chebline(off, scl): ...
+def chebfromroots(roots): ...
+def chebadd(c1, c2): ...
+def chebsub(c1, c2): ...
+def chebmulx(c): ...
+def chebmul(c1, c2): ...
+def chebdiv(c1, c2): ...
+def chebpow(c, pow, maxpower=...): ...
+def chebder(c, m=..., scl=..., axis=...): ...
+def chebint(c, m=..., k = ..., lbnd=..., scl=..., axis=...): ...
+def chebval(x, c, tensor=...): ...
+def chebval2d(x, y, c): ...
+def chebgrid2d(x, y, c): ...
+def chebval3d(x, y, z, c): ...
+def chebgrid3d(x, y, z, c): ...
+def chebvander(x, deg): ...
+def chebvander2d(x, y, deg): ...
+def chebvander3d(x, y, z, deg): ...
+def chebfit(x, y, deg, rcond=..., full=..., w=...): ...
+def chebcompanion(c): ...
+def chebroots(c): ...
+def chebinterpolate(func, deg, args = ...): ...
+def chebgauss(deg): ...
+def chebweight(x): ...
+def chebpts1(npts): ...
+def chebpts2(npts): ...
+
+class Chebyshev(ABCPolyBase):
+    @classmethod
+    def interpolate(cls, func, deg, domain=..., args = ...): ...
+    domain: Any
+    window: Any
+    basis_name: Any
diff --git a/numpy/polynomial/hermite.pyi b/numpy/polynomial/hermite.pyi
new file mode 100644
index 000000000..8364a5b0f
--- /dev/null
+++ b/numpy/polynomial/hermite.pyi
@@ -0,0 +1,46 @@
+from typing import Any, List
+
+from numpy import ndarray, dtype, int_, float_
+from numpy.polynomial._polybase import ABCPolyBase
+from numpy.polynomial.polyutils import trimcoef
+
+__all__: list[str]
+
+hermtrim = trimcoef
+
+def poly2herm(pol): ...
+def herm2poly(c): ...
+
+hermdomain: ndarray[Any, dtype[int_]]
+hermzero: ndarray[Any, dtype[int_]]
+hermone: ndarray[Any, dtype[int_]]
+hermx: ndarray[Any, dtype[float_]]
+
+def hermline(off, scl): ...
+def hermfromroots(roots): ...
+def hermadd(c1, c2): ...
+def hermsub(c1, c2): ...
+def hermmulx(c): ...
+def hermmul(c1, c2): ...
+def hermdiv(c1, c2): ...
+def hermpow(c, pow, maxpower=...): ...
+def hermder(c, m=..., scl=..., axis=...): ...
+def hermint(c, m=..., k = ..., lbnd=..., scl=..., axis=...): ...
+def hermval(x, c, tensor=...): ...
+def hermval2d(x, y, c): ...
+def hermgrid2d(x, y, c): ...
+def hermval3d(x, y, z, c): ...
+def hermgrid3d(x, y, z, c): ...
+def hermvander(x, deg): ...
+def hermvander2d(x, y, deg): ...
+def hermvander3d(x, y, z, deg): ...
+def hermfit(x, y, deg, rcond=..., full=..., w=...): ...
+def hermcompanion(c): ...
+def hermroots(c): ...
+def hermgauss(deg): ...
+def hermweight(x): ...
+
+class Hermite(ABCPolyBase):
+    domain: Any
+    window: Any
+    basis_name: Any
diff --git a/numpy/polynomial/hermite_e.pyi b/numpy/polynomial/hermite_e.pyi
new file mode 100644
index 000000000..c029bfda7
--- /dev/null
+++ b/numpy/polynomial/hermite_e.pyi
@@ -0,0 +1,46 @@
+from typing import Any, List
+
+from numpy import ndarray, dtype, int_
+from numpy.polynomial._polybase import ABCPolyBase
+from numpy.polynomial.polyutils import trimcoef
+
+__all__: list[str]
+
+hermetrim = trimcoef
+
+def poly2herme(pol): ...
+def herme2poly(c): ...
+
+hermedomain: ndarray[Any, dtype[int_]]
+hermezero: ndarray[Any, dtype[int_]]
+hermeone: ndarray[Any, dtype[int_]]
+hermex: ndarray[Any, dtype[int_]]
+
+def hermeline(off, scl): ...
+def hermefromroots(roots): ...
+def hermeadd(c1, c2): ...
+def hermesub(c1, c2): ...
+def hermemulx(c): ...
+def hermemul(c1, c2): ...
+def hermediv(c1, c2): ...
+def hermepow(c, pow, maxpower=...): ...
+def hermeder(c, m=..., scl=..., axis=...): ...
+def hermeint(c, m=..., k = ..., lbnd=..., scl=..., axis=...): ...
+def hermeval(x, c, tensor=...): ...
+def hermeval2d(x, y, c): ...
+def hermegrid2d(x, y, c): ...
+def hermeval3d(x, y, z, c): ...
+def hermegrid3d(x, y, z, c): ...
+def hermevander(x, deg): ...
+def hermevander2d(x, y, deg): ...
+def hermevander3d(x, y, z, deg): ...
+def hermefit(x, y, deg, rcond=..., full=..., w=...): ...
+def hermecompanion(c): ...
+def hermeroots(c): ...
+def hermegauss(deg): ...
+def hermeweight(x): ...
+
+class HermiteE(ABCPolyBase):
+    domain: Any
+    window: Any
+    basis_name: Any
diff --git a/numpy/polynomial/laguerre.pyi b/numpy/polynomial/laguerre.pyi
new file mode 100644
index 000000000..2b9ab34e0
--- /dev/null
+++ b/numpy/polynomial/laguerre.pyi
@@ -0,0 +1,46 @@
+from typing import Any, List
+
+from numpy import ndarray, dtype, int_
+from numpy.polynomial._polybase import ABCPolyBase
+from numpy.polynomial.polyutils import trimcoef
+
+__all__: list[str]
+
+lagtrim = trimcoef
+
+def poly2lag(pol): ...
+def lag2poly(c): ...
+
+lagdomain: ndarray[Any, dtype[int_]]
+lagzero: ndarray[Any, dtype[int_]]
+lagone: ndarray[Any, dtype[int_]]
+lagx: ndarray[Any, dtype[int_]]
+
+def lagline(off, scl): ...
+def lagfromroots(roots): ...
+def lagadd(c1, c2): ...
+def lagsub(c1, c2): ...
+def lagmulx(c): ...
+def lagmul(c1, c2): ...
+def lagdiv(c1, c2): ...
+def lagpow(c, pow, maxpower=...): ...
+def lagder(c, m=..., scl=..., axis=...): ...
+def lagint(c, m=..., k = ..., lbnd=..., scl=..., axis=...): ...
+def lagval(x, c, tensor=...): ...
+def lagval2d(x, y, c): ...
+def laggrid2d(x, y, c): ...
+def lagval3d(x, y, z, c): ...
+def laggrid3d(x, y, z, c): ...
+def lagvander(x, deg): ...
+def lagvander2d(x, y, deg): ...
+def lagvander3d(x, y, z, deg): ...
+def lagfit(x, y, deg, rcond=..., full=..., w=...): ...
+def lagcompanion(c): ...
+def lagroots(c): ...
+def laggauss(deg): ...
+def lagweight(x): ...
+
+class Laguerre(ABCPolyBase):
+    domain: Any
+    window: Any
+    basis_name: Any
diff --git a/numpy/polynomial/legendre.pyi b/numpy/polynomial/legendre.pyi
new file mode 100644
index 000000000..86aef1793
--- /dev/null
+++ b/numpy/polynomial/legendre.pyi
@@ -0,0 +1,46 @@
+from typing import Any, List
+
+from numpy import ndarray, dtype, int_
+from numpy.polynomial._polybase import ABCPolyBase
+from numpy.polynomial.polyutils import trimcoef
+
+__all__: list[str]
+
+legtrim = trimcoef
+
+def poly2leg(pol): ...
+def leg2poly(c): ...
+
+legdomain: ndarray[Any, dtype[int_]]
+legzero: ndarray[Any, dtype[int_]]
+legone: ndarray[Any, dtype[int_]]
+legx: ndarray[Any, dtype[int_]]
+
+def legline(off, scl): ...
+def legfromroots(roots): ...
+def legadd(c1, c2): ...
+def legsub(c1, c2): ...
+def legmulx(c): ...
+def legmul(c1, c2): ...
+def legdiv(c1, c2): ...
+def legpow(c, pow, maxpower=...): ...
+def legder(c, m=..., scl=..., axis=...): ...
+def legint(c, m=..., k = ..., lbnd=..., scl=..., axis=...): ...
+def legval(x, c, tensor=...): ...
+def legval2d(x, y, c): ...
+def leggrid2d(x, y, c): ...
+def legval3d(x, y, z, c): ...
+def leggrid3d(x, y, z, c): ...
+def legvander(x, deg): ...
+def legvander2d(x, y, deg): ...
+def legvander3d(x, y, z, deg): ...
+def legfit(x, y, deg, rcond=..., full=..., w=...): ...
+def legcompanion(c): ...
+def legroots(c): ...
+def leggauss(deg): ...
+def legweight(x): ...
+
+class Legendre(ABCPolyBase):
+    domain: Any
+    window: Any
+    basis_name: Any
diff --git a/numpy/polynomial/polynomial.pyi b/numpy/polynomial/polynomial.pyi
new file mode 100644
index 000000000..f779300a9
--- /dev/null
+++ b/numpy/polynomial/polynomial.pyi
@@ -0,0 +1,41 @@
+from typing import Any, List
+
+from numpy import ndarray, dtype, int_
+from numpy.polynomial._polybase import ABCPolyBase
+from numpy.polynomial.polyutils import trimcoef
+
+__all__: list[str]
+
+polytrim = trimcoef
+
+polydomain: ndarray[Any, dtype[int_]]
+polyzero: ndarray[Any, dtype[int_]]
+polyone: ndarray[Any, dtype[int_]]
+polyx: ndarray[Any, dtype[int_]]
+
+def polyline(off, scl): ...
+def polyfromroots(roots): ...
+def polyadd(c1, c2): ...
+def polysub(c1, c2): ...
+def polymulx(c): ...
+def polymul(c1, c2): ...
+def polydiv(c1, c2): ...
+def polypow(c, pow, maxpower=...): ...
+def polyder(c, m=..., scl=..., axis=...): ...
+def polyint(c, m=..., k=..., lbnd=..., scl=..., axis=...): ...
+def polyval(x, c, tensor=...): ...
+def polyvalfromroots(x, r, tensor=...): ...
+def polyval2d(x, y, c): ...
+def polygrid2d(x, y, c): ...
+def polyval3d(x, y, z, c): ...
+def polygrid3d(x, y, z, c): ...
+def polyvander(x, deg): ...
+def polyvander2d(x, y, deg): ...
+def polyvander3d(x, y, z, deg): ...
+def polyfit(x, y, deg, rcond=..., full=..., w=...): ...
+def polyroots(c): ...
+
+class Polynomial(ABCPolyBase):
+    domain: Any
+    window: Any
+    basis_name: Any
diff --git a/numpy/polynomial/polyutils.pyi b/numpy/polynomial/polyutils.pyi
new file mode 100644
index 000000000..0befa740e
--- /dev/null
+++ b/numpy/polynomial/polyutils.pyi
@@ -0,0 +1,17 @@
+from typing import List
+
+__all__: List[str]
+
+class RankWarning(UserWarning): ...
+class PolyError(Exception): ...
+class PolyDomainError(PolyError): ...
+
+# NOTE: Deprecated
+# class PolyBase: ...
+
+def trimseq(seq): ...
+def as_series(alist, trim=...): ...
+def trimcoef(c, tol=...): ...
+def getdomain(x): ...
+def mapparms(old, new): ...
+def mapdomain(x, old, new): ...