✨ sparse.linalg: complete the matrix norm functions (#300)

scipy-stubs/sparse/linalg/_norm.pyi

@@ -1,5 +1,20 @@
-from scipy._typing import Untyped
+from typing import Literal, TypeAlias, overload
+
+import numpy as np
+import optype as op
+import optype.numpy as onp
+from scipy.sparse._base import _spbase
 
 __all__ = ["norm"]
 
-def norm(x: Untyped, ord: Untyped | None = None, axis: Untyped | None = None) -> Untyped: ...
+_Ord: TypeAlias = Literal["fro", 0, 1, 2, -1] | float
+_Real: TypeAlias = np.int32 | np.int64 | np.float64
+
+###
+
+@overload  # no axis, two axes
+def norm(x: _spbase, ord: _Ord | None = None, axis: tuple[op.CanIndex, op.CanIndex] | None = None) -> _Real: ...
+@overload  # single axis (positional)
+def norm(x: _spbase, ord: _Ord | None, axis: op.CanIndex) -> onp.Array1D[_Real]: ...
+@overload  # single axis (keyword)
+def norm(x: _spbase, ord: _Ord | None = None, *, axis: op.CanIndex) -> onp.Array1D[_Real]: ...

scipy-stubs/sparse/linalg/_onenormest.pyi

@@ -1,14 +1,75 @@
-from scipy._typing import Untyped
+from typing import Literal, TypeAlias, overload
+
+import numpy as np
+import optype.numpy as onp
+from scipy.sparse._base import _spbase
+from ._interface import LinearOperator
 
 __all__ = ["onenormest"]
 
-def onenormest(A: Untyped, t: int = 2, itmax: int = 5, compute_v: bool = False, compute_w: bool = False) -> Untyped: ...
+_Falsy: TypeAlias = Literal[False, 0]
+_Truthy: TypeAlias = Literal[True, 1]
+_Float1D: TypeAlias = onp.Array1D[np.float64]
+_ToMatrix: TypeAlias = onp.ToComplex2D | LinearOperator | _spbase
 
 #
-def sign_round_up(x: Untyped) -> Untyped: ...
-def elementary_vector(n: Untyped, i: Untyped) -> Untyped: ...
-def vectors_are_parallel(v: Untyped, w: Untyped) -> Untyped: ...
-def every_col_of_X_is_parallel_to_a_col_of_Y(X: Untyped, Y: Untyped) -> Untyped: ...
-def column_needs_resampling(i: Untyped, X: Untyped, Y: Untyped | None = None) -> Untyped: ...
-def resample_column(i: Untyped, X: Untyped) -> None: ...
-def less_than_or_close(a: Untyped, b: Untyped) -> Untyped: ...
+
+@overload  # compute_v: falsy, compute_w: falsy
+def onenormest(
+    A: _ToMatrix,
+    t: int = 2,
+    itmax: int = 5,
+    compute_v: _Falsy = False,
+    compute_w: _Falsy = False,
+) -> np.float64: ...
+@overload  # compute_v: falsy, compute_w: truthy  (positional)
+def onenormest(
+    A: _ToMatrix,
+    t: int,
+    itmax: int,
+    compute_v: _Falsy,
+    compute_w: _Truthy,
+) -> tuple[np.float64, _Float1D]: ...
+@overload  # compute_v: falsy, compute_w: truthy  (keyword)
+def onenormest(
+    A: _ToMatrix,
+    t: int = 2,
+    itmax: int = 5,
+    compute_v: _Falsy = False,
+    *,
+    compute_w: _Truthy,
+) -> tuple[np.float64, _Float1D]: ...
+@overload  # compute_v: truthy  (positional), compute_w: falsy
+def onenormest(
+    A: _ToMatrix,
+    t: int,
+    itmax: int,
+    compute_v: _Truthy,
+    compute_w: _Falsy = False,
+) -> tuple[np.float64, _Float1D]: ...
+@overload  # compute_v: truthy  (keyword), compute_w: falsy
+def onenormest(
+    A: _ToMatrix,
+    t: int = 2,
+    itmax: int = 5,
+    *,
+    compute_v: _Truthy,
+    compute_w: _Falsy = False,
+) -> tuple[np.float64, _Float1D]: ...
+@overload  # compute_v: truthy  (positional), compute_w: truthy
+def onenormest(
+    A: _ToMatrix,
+    t: int,
+    itmax: int,
+    compute_v: _Truthy,
+    compute_w: _Truthy,
+) -> tuple[np.float64, _Float1D, _Float1D]: ...
+@overload  # compute_v: truthy  (keyword), compute_w: truthy
+def onenormest(
+    A: _ToMatrix,
+    t: int = 2,
+    itmax: int = 5,
+    *,
+    compute_v: _Truthy,
+    compute_w: _Truthy,
+) -> tuple[np.float64, _Float1D, _Float1D]: ...