✨ sparse: complete identity (#310)

jorenham · Dec 15, 2024 · e3ab140 · e3ab140
1 parent 86db1d6
commit e3ab140
Showing 1 changed file with 86 additions and 65 deletions.
diff --git a/scipy-stubs/sparse/_construct.pyi b/scipy-stubs/sparse/_construct.pyi
@@ -1,9 +1,9 @@
-from typing import Any, TypeAlias, TypeVar, overload
+from typing import Any, Literal, TypeAlias, TypeVar, overload
 
 import numpy as np
 import numpy.typing as npt
 import optype.typing as opt
-from scipy._typing import Untyped
+from scipy._typing import Seed, Untyped, UntypedCallable
 from scipy.sparse import (
     bsr_array,
     bsr_matrix,
@@ -20,7 +20,7 @@ from scipy.sparse import (
     lil_array,
     lil_matrix,
 )
-from ._typing import Scalar, SPFormat, ToDType, ToDTypeBool, ToDTypeComplex, ToDTypeFloat, ToDTypeInt
+from ._typing import Float, Scalar, SPFormat, ToDType, ToDTypeBool, ToDTypeComplex, ToDTypeFloat, ToDTypeInt, ToShape, ToShape2D
 
 __all__ = [
     "block_array",
@@ -63,8 +63,12 @@ _SpArray: TypeAlias = (
     | dok_array[_SCT, _ShapeT]
     | lil_array[_SCT]
 )
+_SpArray1D: TypeAlias = coo_array[_SCT, tuple[int]] | csr_array[_SCT, tuple[int]] | dok_array[_SCT, tuple[int]]
 _SpArray2D: TypeAlias = _SpArray[_SCT, tuple[int, int]]
 
+_FmtDIA: TypeAlias = Literal["dia"]
+_FmtNonDIA: TypeAlias = Literal["bsr", "coo", "csc", "csr", "dok", "lil"]
+
 ###
 
 #
@@ -73,7 +77,7 @@ def spdiags(
     diags: Untyped,
     m: Untyped | None = None,
     n: Untyped | None = None,
-    format: Untyped | None = None,
+    format: SPFormat | None = None,
 ) -> Untyped: ...
 
 #
@@ -82,22 +86,47 @@ def diags_array(
     /,
     *,
     offsets: int = 0,
-    shape: Untyped | None = None,
-    format: Untyped | None = None,
-    dtype: Untyped | None = None,
-) -> Untyped: ...
+    shape: ToShape | None = None,
+    format: SPFormat | None = None,
+    dtype: npt.DTypeLike | None = None,
+) -> _SpArray1D | _SpArray2D: ...
 
 #
 def diags(
     diagonals: Untyped,
     offsets: int = 0,
-    shape: Untyped | None = None,
-    format: Untyped | None = None,
-    dtype: Untyped | None = None,
-) -> Untyped: ...
+    shape: ToShape2D | None = None,
+    format: SPFormat | None = None,
+    dtype: npt.DTypeLike | None = None,
+) -> _SpMatrix: ...
 
 #
-def identity(n: Untyped, dtype: str = "d", format: Untyped | None = None) -> Untyped: ...
+@overload  # dtype like bool, format: None = ...
+def identity(n: opt.AnyInt, dtype: ToDTypeBool, format: _FmtDIA | None = None) -> dia_matrix[np.bool_]: ...
+@overload  # dtype like int, format: None = ...
+def identity(n: opt.AnyInt, dtype: ToDTypeInt, format: _FmtDIA | None = None) -> dia_matrix[np.int_]: ...
+@overload  # dtype like float (default), format: None = ...
+def identity(n: opt.AnyInt, dtype: ToDTypeFloat = "d", format: _FmtDIA | None = None) -> dia_matrix[np.float64]: ...
+@overload  # dtype like complex, format: None = ...
+def identity(n: opt.AnyInt, dtype: ToDTypeComplex, format: _FmtDIA | None = None) -> dia_matrix[np.complex128]: ...
+@overload  # dtype like <known>, format: None = ...
+def identity(n: opt.AnyInt, dtype: ToDType[_SCT], format: _FmtDIA | None = None) -> dia_matrix[_SCT]: ...
+@overload  # dtype like <unknown>, format: None = ...
+def identity(n: opt.AnyInt, dtype: npt.DTypeLike, format: _FmtDIA | None = None) -> dia_matrix: ...
+@overload  # dtype like float, format: <given> (positional)
+def identity(n: opt.AnyInt, dtype: ToDTypeFloat, format: _FmtNonDIA) -> _SpMatrix[np.float64]: ...
+@overload  # dtype like float (default), format: <given> (keyword)
+def identity(n: opt.AnyInt, dtype: ToDTypeFloat = "d", *, format: _FmtNonDIA) -> _SpMatrix[np.float64]: ...
+@overload  # dtype like bool, format: <given>
+def identity(n: opt.AnyInt, dtype: ToDTypeBool, format: _FmtNonDIA) -> _SpMatrix[np.bool_]: ...
+@overload  # dtype like int, format: <given>
+def identity(n: opt.AnyInt, dtype: ToDTypeInt, format: _FmtNonDIA) -> _SpMatrix[np.int_]: ...
+@overload  # dtype like complex, format: <given>
+def identity(n: opt.AnyInt, dtype: ToDTypeComplex, format: _FmtNonDIA) -> _SpMatrix[np.complex128]: ...
+@overload  # dtype like <known>, format: <given>
+def identity(n: opt.AnyInt, dtype: ToDType[_SCT], format: _FmtNonDIA) -> _SpMatrix[_SCT]: ...
+@overload  # dtype like <unknown>, fformat: <given>
+def identity(n: opt.AnyInt, dtype: npt.DTypeLike, format: _FmtNonDIA) -> _SpMatrix: ...
 
 #
 @overload  # dtype like bool, format: None = ...
@@ -107,7 +136,7 @@ def eye_array(
     *,
     k: int = 0,
     dtype: ToDTypeBool,
-    format: None = None,
+    format: _FmtDIA | None = None,
 ) -> dia_array[np.bool_]: ...
 @overload  # dtype like int, format: None = ...
 def eye_array(
@@ -116,7 +145,7 @@ def eye_array(
     *,
     k: int = 0,
     dtype: ToDTypeInt,
-    format: None = None,
+    format: _FmtDIA | None = None,
 ) -> dia_array[np.int_]: ...
 @overload  # dtype like float (default), format: None = ...
 def eye_array(
@@ -125,7 +154,7 @@ def eye_array(
     *,
     k: int = 0,
     dtype: ToDTypeFloat = ...,
-    format: None = None,
+    format: _FmtDIA | None = None,
 ) -> dia_array[np.float64]: ...
 @overload  # dtype like complex, format: None = ...
 def eye_array(
@@ -134,7 +163,7 @@ def eye_array(
     *,
     k: int = 0,
     dtype: ToDTypeComplex,
-    format: None = None,
+    format: _FmtDIA | None = None,
 ) -> dia_array[np.complex128]: ...
 @overload  # dtype like <known>, format: None = ...
 def eye_array(
@@ -143,7 +172,7 @@ def eye_array(
     *,
     k: int = 0,
     dtype: ToDType[_SCT],
-    format: None = None,
+    format: _FmtDIA | None = None,
 ) -> dia_array[_SCT]: ...
 @overload  # dtype like <unknown>, format: None = ...
 def eye_array(
@@ -152,7 +181,7 @@ def eye_array(
     *,
     k: int = 0,
     dtype: npt.DTypeLike,
-    format: None = None,
+    format: _FmtDIA | None = None,
 ) -> dia_array: ...
 @overload  # dtype like float (default), format: <given>
 def eye_array(
@@ -161,7 +190,7 @@ def eye_array(
     *,
     k: int = 0,
     dtype: ToDTypeFloat = ...,
-    format: SPFormat,
+    format: _FmtNonDIA,
 ) -> _SpArray2D[np.float64]: ...
 @overload  # dtype like bool, format: <given>
 def eye_array(
@@ -170,7 +199,7 @@ def eye_array(
     *,
     k: int = 0,
     dtype: ToDTypeBool,
-    format: SPFormat,
+    format: _FmtNonDIA,
 ) -> _SpArray2D[np.bool_]: ...
 @overload  # dtype like int, format: <given>
 def eye_array(
@@ -179,7 +208,7 @@ def eye_array(
     *,
     k: int = 0,
     dtype: ToDTypeInt,
-    format: SPFormat,
+    format: _FmtNonDIA,
 ) -> _SpArray2D[np.int_]: ...
 @overload  # dtype like complex, format: <given>
 def eye_array(
@@ -188,7 +217,7 @@ def eye_array(
     *,
     k: int = 0,
     dtype: ToDTypeComplex,
-    format: SPFormat,
+    format: _FmtNonDIA,
 ) -> _SpArray2D[np.complex128]: ...
 @overload  # dtype like <known>, format: <given>
 def eye_array(
@@ -197,7 +226,7 @@ def eye_array(
     *,
     k: int = 0,
     dtype: ToDType[_SCT],
-    format: SPFormat,
+    format: _FmtNonDIA,
 ) -> _SpArray2D[_SCT]: ...
 @overload  # dtype like <unknown>, fformat: <given>
 def eye_array(
@@ -206,10 +235,10 @@ def eye_array(
     *,
     k: int = 0,
     dtype: npt.DTypeLike,
-    format: SPFormat,
+    format: _FmtNonDIA,
 ) -> _SpArray2D: ...
 
-#
+# NOTE: `eye_array` should be prefered over `eye`
 @overload  # dtype like float (default)
 def eye(
     m: opt.AnyInt,
@@ -305,54 +334,46 @@ def eye(
 ) -> _SpMatrix: ...
 
 #
-def kron(A: Untyped, B: Untyped, format: Untyped | None = None) -> Untyped: ...
-
-#
-def kronsum(A: Untyped, B: Untyped, format: Untyped | None = None) -> Untyped: ...
-
-#
-def hstack(blocks: Untyped, format: Untyped | None = None, dtype: Untyped | None = None) -> Untyped: ...
-
-#
-def vstack(blocks: Untyped, format: Untyped | None = None, dtype: Untyped | None = None) -> Untyped: ...
-
-#
-def bmat(blocks: Untyped, format: Untyped | None = None, dtype: Untyped | None = None) -> Untyped: ...
+def kron(A: Untyped, B: Untyped, format: SPFormat | None = None) -> _SpArray2D | _SpMatrix: ...
+def kronsum(A: Untyped, B: Untyped, format: SPFormat | None = None) -> _SpArray2D | _SpMatrix: ...
 
 #
-def block_array(blocks: Untyped, *, format: Untyped | None = None, dtype: Untyped | None = None) -> Untyped: ...
+def hstack(blocks: Untyped, format: SPFormat | None = None, dtype: npt.DTypeLike | None = None) -> _SpArray | _SpMatrix: ...
+def vstack(blocks: Untyped, format: SPFormat | None = None, dtype: npt.DTypeLike | None = None) -> _SpArray | _SpMatrix: ...
 
 #
-def block_diag(mats: Untyped, format: Untyped | None = None, dtype: Untyped | None = None) -> Untyped: ...
+def bmat(blocks: Untyped, format: SPFormat | None = None, dtype: npt.DTypeLike | None = None) -> _SpArray | _SpMatrix: ...
+def block_array(blocks: Untyped, *, format: SPFormat | None = None, dtype: npt.DTypeLike | None = None) -> _SpArray: ...
+def block_diag(mats: Untyped, format: SPFormat | None = None, dtype: npt.DTypeLike | None = None) -> _SpArray | _SpMatrix: ...
 
 #
 def random_array(
-    shape: Untyped,
+    shape: ToShape,
     *,
-    density: float = 0.01,
-    format: str = "coo",
-    dtype: Untyped | None = None,
-    random_state: Untyped | None = None,
-    data_sampler: Untyped | None = None,
-) -> Untyped: ...
+    density: float | Float = 0.01,
+    format: SPFormat = "coo",
+    dtype: npt.DTypeLike | None = None,
+    random_state: Seed | None = None,
+    data_sampler: UntypedCallable | None = None,
+) -> _SpArray1D | _SpArray2D: ...
 
-#
+# NOTE: `random_array` should be prefered over `random`
 def random(
-    m: Untyped,
-    n: Untyped,
-    density: float = 0.01,
-    format: str = "coo",
-    dtype: Untyped | None = None,
-    random_state: Untyped | None = None,
-    data_rvs: Untyped | None = None,
-) -> Untyped: ...
+    m: opt.AnyInt,
+    n: opt.AnyInt,
+    density: float | Float = 0.01,
+    format: SPFormat = "coo",
+    dtype: npt.DTypeLike | None = None,
+    random_state: Seed | None = None,
+    data_rvs: UntypedCallable | None = None,
+) -> _SpMatrix: ...
 
-#
+# NOTE: `random_array` should be prefered over `rand`
 def rand(
-    m: Untyped,
-    n: Untyped,
-    density: float = 0.01,
-    format: str = "coo",
-    dtype: Untyped | None = None,
-    random_state: Untyped | None = None,
-) -> Untyped: ...
+    m: opt.AnyInt,
+    n: opt.AnyInt,
+    density: float | Float = 0.01,
+    format: SPFormat = "coo",
+    dtype: npt.DTypeLike | None = None,
+    random_state: Seed | None = None,
+) -> _SpMatrix: ...