✨ complete interpolate._fitpack* (1/2) (#275)

scipy-stubs/interpolate/_dfitpack.pyi

@@ -1,10 +1,12 @@
 from collections.abc import Callable
-from typing import Any, Final, Literal as L, Protocol, final, type_check_only
+from typing import Any, Final, Literal as L, Protocol, TypeAlias, final, type_check_only
 from typing_extensions import LiteralString, TypeVar
 
 import numpy as np
 import optype.numpy as onp
 
+_Float1D: TypeAlias = onp.Array1D[np.float64]
+
 _NameT = TypeVar("_NameT", bound=str)
 _FuncT_co = TypeVar("_FuncT_co", bound=Callable[..., object], default=Callable[..., Any], covariant=True)
 
@@ -34,7 +36,7 @@ class _Func_splev(Protocol):
         k: onp.ToJustInt,  # 0 <= k < len(c)
         x: onp.ToFloat1D,
         e: L[0, 1, 2, 3] = 0,
-    ) -> tuple[onp.Array1D[np.float64], int]: ...
+    ) -> tuple[_Float1D, int]: ...
 
 @type_check_only
 class _Func_splder(Protocol):
@@ -47,7 +49,7 @@ class _Func_splder(Protocol):
         x: onp.ToFloat1D,
         nu: onp.ToJustInt = 1,  # 0 <= nu <= k < len(c)
         e: L[0, 1, 2, 3] = 0,
-    ) -> tuple[onp.Array1D[np.float64], int]: ...
+    ) -> tuple[_Float1D, int]: ...
 
 @type_check_only
 class _Func_splint(Protocol):
@@ -59,7 +61,7 @@ class _Func_splint(Protocol):
         k: onp.ToJustInt,  # 0 <= k < len(c)
         a: onp.ToFloat | onp.ToFloatND,
         b: onp.ToFloat | onp.ToFloatND,
-    ) -> tuple[float, onp.Array1D[np.float64]]: ...
+    ) -> tuple[float, _Float1D]: ...
 
 @type_check_only
 class _Func_sproot(Protocol):
@@ -69,7 +71,7 @@ class _Func_sproot(Protocol):
         t: onp.ToFloat1D,  # len(t) >= 8
         c: onp.ToFloat1D,  # len(c) >= len(t) - 4
         mest: onp.ToJustInt = ...,  # mest = 3 * (len(t) - 7)
-    ) -> tuple[onp.Array1D[np.float64], int, int]: ...
+    ) -> tuple[_Float1D, int, int]: ...
 
 @type_check_only
 class _Func_spalde(Protocol):
@@ -80,7 +82,7 @@ class _Func_spalde(Protocol):
         c: onp.ToFloat1D,  # len(c) >= len(t) - 4
         k1: onp.ToJustInt,
         x: onp.ToFloat | onp.ToFloatND,
-    ) -> tuple[onp.Array1D[np.float64], int]: ...
+    ) -> tuple[_Float1D, int]: ...
 
 @type_check_only
 class _Func_bispe_(Protocol):
@@ -94,7 +96,7 @@ class _Func_bispe_(Protocol):
         ky: onp.ToJustInt,  # 0 <= ky < len(y)
         x: onp.ToFloat1D,
         y: onp.ToFloat1D,
-    ) -> tuple[onp.Array1D[np.float64], int]: ...
+    ) -> tuple[_Float1D, int]: ...
 
 ###
 
@@ -116,13 +118,17 @@ splint: _FortranFunction[L["splint"], _Func_splint] = ...
 sproot: _FortranFunction[L["function sproot"], _Func_sproot] = ...
 # (t, c, k1, x) -> (d, ier)
 spalde: _FortranFunction[L["function spalde"], _Func_spalde] = ...
-# TODO(jorenham)
-# (iopt, x, y, w, t, wrk, iwrk, *, xb=x[0], xe=x[-1], k=3, s=0.0) -> (t, c, fp, wrk, iwrk, ier)
-curfit: _FortranFunction[L["function curfit"]] = ...
-# (iopt, x, y, w, t, wrk, iwrk, *, k=3, s=0.0) -> (t, c, fp, wrk, iwrk, ier)
-percur: _FortranFunction[L["function percur"]] = ...
-# (iopt, ipar, idim, u, x, w, ub, ue, t, wrk, iwrk, *, k=3, s=0.0) -> (t, c, fp, wrk, iwrk, ier)
-parcur: _FortranFunction[L["function parcur"]] = ...
+
+_Func_cur: TypeAlias = Callable[..., tuple[int, _Float1D, float, int]]
+
+# (iopt, x, y, w, t, wrk, iwrk, xb?, xe?, k?, s?) -> (n, c, fp, ier)
+curfit: _FortranFunction[L["function curfit"], _Func_cur] = ...
+# (iopt, x, y, w, t, wrk, iwrk, k?, s?) -> (n, c, fp, ier)
+percur: _FortranFunction[L["function percur"], _Func_cur] = ...
+# (iopt, ipar, idim, u, x, w, ub, ue, t, wrk, iwrk, k?, s?) -> (n, c, fp, ier)
+parcur: _FortranFunction[L["function parcur"], _Func_cur] = ...
+
+# these have ridculously large signatures...
 fpcurf0: _FortranFunction[L["function fpcurf0"]] = ...
 fpcurf1: _FortranFunction[L["function fpcurf1"]] = ...
 fpcurfm1: _FortranFunction[L["function fpcurfm1"]] = ...

scipy-stubs/interpolate/_fitpack_impl.pyi

@@ -1,42 +1,152 @@
+from collections.abc import Sequence
+from typing import Literal, TypeAlias, overload
+from typing_extensions import LiteralString
+
+import numpy as np
+import optype.numpy as onp
 from scipy._typing import Untyped
 
 __all__ = ["bisplev", "bisplrep", "insert", "spalde", "splantider", "splder", "splev", "splint", "splprep", "splrep", "sproot"]
 
+_Falsy: TypeAlias = Literal[False, 0]
+_Truthy: TypeAlias = Literal[True, 1]
+
+_Float: TypeAlias = float | np.float64
+_Float1D: TypeAlias = onp.Array1D[np.float64]
+_FloatND: TypeAlias = onp.ArrayND[np.float64]
+
+_Task: TypeAlias = Literal[-1, 0, 1]
+_Ext: TypeAlias = Literal[0, 1, 2, 3]
+
+_ToTCK: TypeAlias = Sequence[onp.ToFloat1D | onp.ToFloat2D | int]
+# `(t, c, k)`
+_OutTCK: TypeAlias = tuple[_Float1D, _Float1D, int]
+# `([t, c, k], u)`
+_OutTCKU: TypeAlias = tuple[Sequence[_Float1D | list[_Float1D] | int], _Float1D]
+
+###
+
+# NOTE: The docs are incorrect about the return type of `splgrep`
+@overload  # full_output: falsy = ...
 def splprep(
-    x: Untyped,
-    w: Untyped | None = None,
-    u: Untyped | None = None,
-    ub: Untyped | None = None,
-    ue: Untyped | None = None,
+    x: onp.ToFloat2D,
+    w: onp.ToFloat1D | None = None,
+    u: onp.ToFloat1D | None = None,
+    ub: onp.ToFloat | None = None,
+    ue: onp.ToFloat | None = None,
     k: int = 3,
-    task: int = 0,
-    s: Untyped | None = None,
-    t: Untyped | None = None,
-    full_output: int = 0,
-    nest: Untyped | None = None,
-    per: int = 0,
-    quiet: int = 1,
-) -> Untyped: ...
+    task: _Task = 0,
+    s: onp.ToFloat | None = None,
+    t: onp.ToFloat1D | None = None,
+    full_output: _Falsy = 0,
+    nest: int | None = None,
+    per: onp.ToBool = 0,
+    quiet: onp.ToBool = 1,
+) -> _OutTCKU: ...
+@overload  # full_output: truthy (positional)
+def splprep(
+    x: onp.ToFloat2D,
+    w: onp.ToFloat1D | None,
+    u: onp.ToFloat1D | None,
+    ub: onp.ToFloat | None,
+    ue: onp.ToFloat | None,
+    k: int,
+    task: _Task,
+    s: onp.ToFloat | None,
+    t: onp.ToFloat1D | None,
+    full_output: _Truthy,
+    nest: int | None = None,
+    per: onp.ToBool = 0,
+    quiet: onp.ToBool = 1,
+) -> tuple[_OutTCKU, _Float, int, LiteralString]: ...
+@overload  # full_output: truthy (keyword)
+def splprep(
+    x: onp.ToFloat2D,
+    w: onp.ToFloat1D | None = None,
+    u: onp.ToFloat1D | None = None,
+    ub: onp.ToFloat | None = None,
+    ue: onp.ToFloat | None = None,
+    k: int = 3,
+    task: _Task = 0,
+    s: onp.ToFloat | None = None,
+    t: onp.ToFloat1D | None = None,
+    *,
+    full_output: _Truthy,
+    nest: int | None = None,
+    per: onp.ToBool = 0,
+    quiet: onp.ToBool = 1,
+) -> tuple[_OutTCKU, _Float, int, LiteralString]: ...
+
+#
+@overload  # full_output: falsy = ...
 def splrep(
-    x: Untyped,
-    y: Untyped,
-    w: Untyped | None = None,
-    xb: Untyped | None = None,
-    xe: Untyped | None = None,
+    x: onp.ToFloat1D,
+    y: onp.ToFloat1D,
+    w: onp.ToFloat1D | None = None,
+    xb: onp.ToFloat | None = None,
+    xe: onp.ToFloat | None = None,
     k: int = 3,
-    task: int = 0,
-    s: Untyped | None = None,
-    t: Untyped | None = None,
-    full_output: int = 0,
-    per: int = 0,
-    quiet: int = 1,
-) -> Untyped: ...
-def splev(x: Untyped, tck: Untyped, der: int = 0, ext: int = 0) -> Untyped: ...
-def splint(a: Untyped, b: Untyped, tck: Untyped, full_output: int = 0) -> Untyped: ...
-def sproot(tck: Untyped, mest: int = 10) -> Untyped: ...
-def spalde(x: Untyped, tck: Untyped) -> Untyped: ...
+    task: _Task = 0,
+    s: onp.ToFloat | None = None,
+    t: onp.ToFloat1D | None = None,
+    full_output: _Falsy = 0,
+    per: onp.ToBool = 0,
+    quiet: onp.ToBool = 1,
+) -> _OutTCK: ...
+@overload  # full_output: truthy (positional)
+def splrep(
+    x: onp.ToFloat1D,
+    y: onp.ToFloat1D,
+    w: onp.ToFloat1D | None,
+    xb: onp.ToFloat | None,
+    xe: onp.ToFloat | None,
+    k: int,
+    task: _Task,
+    s: onp.ToFloat | None,
+    t: onp.ToFloat1D | None,
+    full_output: _Truthy,
+    per: onp.ToBool = 0,
+    quiet: onp.ToBool = 1,
+) -> tuple[_OutTCK, _Float, int, LiteralString]: ...
+@overload  # full_output: truthy (keyword)
+def splrep(
+    x: onp.ToFloat1D,
+    y: onp.ToFloat1D,
+    w: onp.ToFloat1D | None = None,
+    xb: onp.ToFloat | None = None,
+    xe: onp.ToFloat | None = None,
+    k: int = 3,
+    task: _Task = 0,
+    s: onp.ToFloat | None = None,
+    t: onp.ToFloat1D | None = None,
+    *,
+    full_output: _Truthy,
+    per: onp.ToBool = 0,
+    quiet: onp.ToBool = 1,
+) -> tuple[_OutTCK, _Float, int, LiteralString]: ...
+
+#
+def splev(x: onp.ToFloatND, tck: _ToTCK, der: int = 0, ext: _Ext = 0) -> _FloatND: ...
+
+#
+@overload  # full_output: falsy
+def splint(a: onp.ToFloat, b: onp.ToFloat, tck: _ToTCK, full_output: _Falsy = 0) -> _Float | list[_Float]: ...
+@overload  # full_output: truthy
+def splint(a: onp.ToFloat, b: onp.ToFloat, tck: _ToTCK, full_output: _Truthy) -> tuple[_Float | list[_Float], _Float1D]: ...
+
+#
+def sproot(tck: _ToTCK, mest: int = 10) -> _Float1D | list[_Float1D]: ...
+
+#
+@overload
+def spalde(x: onp.ToFloatStrict1D, tck: _ToTCK) -> _Float1D: ...
+@overload
+def spalde(x: onp.ToFloatStrict2D, tck: _ToTCK) -> list[_Float1D]: ...
+@overload
+def spalde(x: onp.ToFloat1D | onp.ToFloat2D, tck: _ToTCK) -> _Float1D | list[_Float1D]: ...
 
 #
+# TODO(jorenham): full_output=True
 def bisplrep(
     x: Untyped,
     y: Untyped,
@@ -48,18 +158,28 @@ def bisplrep(
     ye: Untyped | None = None,
     kx: int = 3,
     ky: int = 3,
-    task: int = 0,
+    task: _Task = 0,
     s: Untyped | None = None,
     eps: float = 1e-16,
     tx: Untyped | None = None,
     ty: Untyped | None = None,
-    full_output: int = 0,
+    full_output: _Falsy = 0,
     nxest: Untyped | None = None,
     nyest: Untyped | None = None,
     quiet: int = 1,
 ) -> Untyped: ...
+
+#
 def bisplev(x: Untyped, y: Untyped, tck: Untyped, dx: int = 0, dy: int = 0) -> Untyped: ...
+
+#
 def dblint(xa: Untyped, xb: Untyped, ya: Untyped, yb: Untyped, tck: Untyped) -> Untyped: ...
+
+#
 def insert(x: Untyped, tck: Untyped, m: int = 1, per: int = 0) -> Untyped: ...
+
+#
 def splder(tck: Untyped, n: int = 1) -> Untyped: ...
+
+#
 def splantider(tck: Untyped, n: int = 1) -> Untyped: ...

scipy-stubs/interpolate/_fitpack_py.pyi

@@ -1,3 +1,58 @@
-from ._fitpack_impl import bisplev, bisplrep, insert, spalde, splantider, splder, splev, splint, splprep, splrep, sproot
+from collections.abc import Sequence
+from typing import Literal, TypeAlias, overload
+
+import numpy as np
+import optype.numpy as onp
+from ._bsplines import BSpline
+from ._fitpack_impl import bisplev, bisplrep, insert, splantider, splder, splprep, splrep
 
 __all__ = ["bisplev", "bisplrep", "insert", "spalde", "splantider", "splder", "splev", "splint", "splprep", "splrep", "sproot"]
+
+_Falsy: TypeAlias = Literal[False, 0]
+_Truthy: TypeAlias = Literal[True, 1]
+
+_Float: TypeAlias = float | np.float64
+_Float1D: TypeAlias = onp.Array1D[np.float64]
+_Float2D: TypeAlias = onp.Array2D[np.float64]
+_FloatND: TypeAlias = onp.ArrayND[np.float64]
+
+_Ext: TypeAlias = Literal[0, 1, 2, 3]
+_ToTCK: TypeAlias = Sequence[onp.ToFloat1D | onp.ToFloat2D | int]
+
+###
+
+# NOTE: The following functions also accept `BSpline` instances, unlike their duals in `_fitpack_impl`.
+
+#
+@overload  # tck: BSpline
+def splev(x: onp.ToFloatND, tck: BSpline, der: int = 0, ext: _Ext = 0) -> _FloatND: ...
+@overload  # tck: (t, c, k)
+def splev(x: onp.ToFloatND, tck: _ToTCK, der: int = 0, ext: _Ext = 0) -> _FloatND | list[_FloatND]: ...
+
+#
+@overload  # tck: BSpline, full_output: falsy
+def splint(a: onp.ToFloat, b: onp.ToFloat, tck: BSpline, full_output: _Falsy = 0) -> _Float | _Float1D: ...
+@overload  # tck: BSpline, full_output: truthy
+def splint(a: onp.ToFloat, b: onp.ToFloat, tck: BSpline, full_output: _Truthy) -> tuple[_Float | _Float1D, _Float1D]: ...
+@overload  # tck: (t, c, k), full_output: falsy
+def splint(a: onp.ToFloat, b: onp.ToFloat, tck: _ToTCK, full_output: _Falsy = 0) -> _Float | list[_Float]: ...
+@overload  # tck: (t, c, k), full_output: truthy
+def splint(a: onp.ToFloat, b: onp.ToFloat, tck: _ToTCK, full_output: _Truthy) -> tuple[_Float | list[_Float], _Float1D]: ...
+
+#
+@overload  # tck: BSpline
+def sproot(tck: BSpline, mest: int = 10) -> _Float1D | _Float2D: ...
+@overload  # tck: (t, c, k)
+def sproot(tck: _ToTCK, mest: int = 10) -> _Float1D | list[_Float1D]: ...
+
+#
+@overload  # x: 1-d
+def spalde(x: onp.ToFloatStrict1D, tck: BSpline | _ToTCK) -> _Float1D: ...
+@overload  # x: 2-d, tck: BSpline
+def spalde(x: onp.ToFloatStrict2D, tck: BSpline) -> _Float2D: ...
+@overload  # x: 2-d, tck: (t, c, k)
+def spalde(x: onp.ToFloatStrict2D, tck: _ToTCK) -> list[_Float1D]: ...
+@overload  # tck: BSpline
+def spalde(x: onp.ToFloat1D | onp.ToFloat2D, tck: BSpline) -> _Float1D | _Float2D: ...
+@overload  # tck: (t, c, k)
+def spalde(x: onp.ToFloat1D | onp.ToFloat2D, tck: _ToTCK) -> _Float1D | list[_Float1D]: ...