diff --git a/scipy-stubs/sparse/_construct.pyi b/scipy-stubs/sparse/_construct.pyi index 9c86dba7..460a5994 100644 --- a/scipy-stubs/sparse/_construct.pyi +++ b/scipy-stubs/sparse/_construct.pyi @@ -3,7 +3,6 @@ from collections.abc import Callable, Iterable, Sequence as Seq from typing import Any, Literal, Never, Protocol, TypeAlias, TypeVar, overload, type_check_only import numpy as np -import numpy.typing as npt import optype.numpy as onp import optype.numpy.compat as npc import optype.typing as opt @@ -45,6 +44,8 @@ _ShapeT = TypeVar("_ShapeT", bound=tuple[int, *tuple[int, ...]]) _Numeric: TypeAlias = npc.number | np.bool_ +_ToDType: TypeAlias = type[complex | _Numeric] | np.dtype[_Numeric] | str + _SpBase2D: TypeAlias = _spbase[_SCT, tuple[int, int]] _SpArray2D: TypeAlias = sparray[_SCT, tuple[int, int]] @@ -56,7 +57,7 @@ _DIAArray: TypeAlias = dia_array[_SCT] _DOKArray2D: TypeAlias = dok_array[_SCT, tuple[int, int]] _LILArray: TypeAlias = lil_array[_SCT] -_ToArray1D2D: TypeAlias = onp.ToArray1D[Never, _SCT] | onp.ToArray2D[Never, _SCT] +_ToArray1D2D: TypeAlias = onp.CanArray[tuple[int] | tuple[int, int], np.dtype[_SCT]] | Seq[_SCT | Seq[_SCT]] _ToSpMatrix: TypeAlias = spmatrix[_SCT] | onp.ToArray2D[Never, _SCT] _ToSparse2D: TypeAlias = _SpBase2D[_SCT] | onp.ToArray2D[Never, _SCT] @@ -76,7 +77,16 @@ _Offsets: TypeAlias = onp.ToInt | onp.ToInt1D _DataRVS: TypeAlias = Callable[[int], onp.ArrayND[_Numeric]] -_ToBlocks: TypeAlias = Seq[Seq[_spbase[_SCT] | None]] | onp.ArrayND[np.object_] +_ToBlocks: TypeAlias = Seq[Seq[_T | None]] +_ToBlocksSPArray: TypeAlias = _ToBlocks[_SpArray2D[_SCT]] +_ToBlocksSPMatrix: TypeAlias = _ToBlocks[spmatrix[_SCT]] +_ToBlocksCanStack: TypeAlias = _ToBlocks[_CanStack[_T]] +_ToBlocksCanStackAs: TypeAlias = _ToBlocks[_CanStackAs[_SCT0, _T]] +_ToBlocksArray: TypeAlias = onp.ArrayND[np.object_] +_ToBlocksUnkown: TypeAlias = _ToBlocksSPArray | _ToBlocksArray + +_ToMatsDiag: TypeAlias = Iterable[_spbase[_SCT] | onp.ToArrayND[_SCT]] +_ToMatsDiagUnknown: TypeAlias = Iterable[_spbase | onp.ArrayND[_Numeric] | complex | Seq[onp.ToComplex] | Seq[onp.ToComplex1D]] @type_check_only class _DataSampler(Protocol): @@ -559,31 +569,31 @@ def diags_array( offsets: _Offsets = 0, shape: _ToShape2D | None = None, format: _FmtDIA | None = None, - dtype: npt.DTypeLike, + dtype: _ToDType, ) -> _DIAArray: ... @overload # diagonals: , format: "bsr", dtype: def diags_array( - diagonals: _ToComplex1D2D, /, *, offsets: _Offsets = 0, shape: _ToShape2D | None = None, format: _FmtBSR, dtype: npt.DTypeLike + diagonals: _ToComplex1D2D, /, *, offsets: _Offsets = 0, shape: _ToShape2D | None = None, format: _FmtBSR, dtype: _ToDType ) -> _BSRArray: ... @overload # diagonals: , format: "coo", dtype: def diags_array( - diagonals: _ToComplex1D2D, /, *, offsets: _Offsets = 0, shape: _ToShape2D | None = None, format: _FmtCOO, dtype: npt.DTypeLike + diagonals: _ToComplex1D2D, /, *, offsets: _Offsets = 0, shape: _ToShape2D | None = None, format: _FmtCOO, dtype: _ToDType ) -> _COOArray2D: ... @overload # diagonals: , format: "csc", dtype: def diags_array( - diagonals: _ToComplex1D2D, /, *, offsets: _Offsets = 0, shape: _ToShape2D | None = None, format: _FmtCSC, dtype: npt.DTypeLike + diagonals: _ToComplex1D2D, /, *, offsets: _Offsets = 0, shape: _ToShape2D | None = None, format: _FmtCSC, dtype: _ToDType ) -> _CSCArray: ... @overload # diagonals: , format: "csr", dtype: def diags_array( - diagonals: _ToComplex1D2D, /, *, offsets: _Offsets = 0, shape: _ToShape2D | None = None, format: _FmtCSR, dtype: npt.DTypeLike + diagonals: _ToComplex1D2D, /, *, offsets: _Offsets = 0, shape: _ToShape2D | None = None, format: _FmtCSR, dtype: _ToDType ) -> _CSRArray2D: ... @overload # diagonals: , format: "dok", dtype: def diags_array( - diagonals: _ToComplex1D2D, /, *, offsets: _Offsets = 0, shape: _ToShape2D | None = None, format: _FmtDOK, dtype: npt.DTypeLike + diagonals: _ToComplex1D2D, /, *, offsets: _Offsets = 0, shape: _ToShape2D | None = None, format: _FmtDOK, dtype: _ToDType ) -> _DOKArray2D: ... @overload # diagonals: , format: "lil", dtype: def diags_array( - diagonals: _ToComplex1D2D, /, *, offsets: _Offsets = 0, shape: _ToShape2D | None = None, format: _FmtLIL, dtype: npt.DTypeLike + diagonals: _ToComplex1D2D, /, *, offsets: _Offsets = 0, shape: _ToShape2D | None = None, format: _FmtLIL, dtype: _ToDType ) -> _LILArray: ... ### @@ -663,7 +673,7 @@ def diags( offsets: _Offsets = 0, shape: _ToShape2D | None = None, format: _FmtDIA | None = None, - dtype: npt.DTypeLike | None = None, + dtype: _ToDType | None = None, ) -> dia_matrix: ... @overload # diagonals: , format: "bsr", dtype: def diags( @@ -672,7 +682,7 @@ def diags( shape: _ToShape2D | None = None, *, format: _FmtBSR, - dtype: npt.DTypeLike | None = None, + dtype: _ToDType | None = None, ) -> bsr_matrix: ... @overload # diagonals: , format: "coo", dtype: def diags( @@ -681,7 +691,7 @@ def diags( shape: _ToShape2D | None = None, *, format: _FmtCOO, - dtype: npt.DTypeLike | None = None, + dtype: _ToDType | None = None, ) -> coo_matrix: ... @overload # diagonals: , format: "csr", dtype: def diags( @@ -690,7 +700,7 @@ def diags( shape: _ToShape2D | None = None, *, format: _FmtCSR, - dtype: npt.DTypeLike | None = None, + dtype: _ToDType | None = None, ) -> csr_matrix: ... @overload # diagonals: , format: "csc", dtype: def diags( @@ -699,7 +709,7 @@ def diags( shape: _ToShape2D | None = None, *, format: _FmtCSC, - dtype: npt.DTypeLike | None = None, + dtype: _ToDType | None = None, ) -> csc_matrix: ... @overload # diagonals: , format: "dok", dtype: def diags( @@ -708,7 +718,7 @@ def diags( shape: _ToShape2D | None = None, *, format: _FmtDOK, - dtype: npt.DTypeLike | None = None, + dtype: _ToDType | None = None, ) -> dok_matrix: ... @overload # diagonals: , format: "lil", dtype: def diags( @@ -717,7 +727,7 @@ def diags( shape: _ToShape2D | None = None, *, format: _FmtLIL, - dtype: npt.DTypeLike | None = None, + dtype: _ToDType | None = None, ) -> lil_matrix: ... ### @@ -897,19 +907,19 @@ def identity(n: opt.AnyInt, dtype: onp.ToDType[_SCT], format: _FmtLIL) -> lil_ma # @overload # dtype: , format: "dia" | None -def identity(n: opt.AnyInt, dtype: npt.DTypeLike, format: _FmtDIA | None = None) -> dia_matrix[Incomplete]: ... +def identity(n: opt.AnyInt, dtype: _ToDType, format: _FmtDIA | None = None) -> dia_matrix[Incomplete]: ... @overload # dtype: , format: "bsr" -def identity(n: opt.AnyInt, dtype: npt.DTypeLike, format: _FmtBSR) -> bsr_matrix[Incomplete]: ... +def identity(n: opt.AnyInt, dtype: _ToDType, format: _FmtBSR) -> bsr_matrix[Incomplete]: ... @overload # dtype: , format: "coo" -def identity(n: opt.AnyInt, dtype: npt.DTypeLike, format: _FmtCOO) -> coo_matrix[Incomplete]: ... +def identity(n: opt.AnyInt, dtype: _ToDType, format: _FmtCOO) -> coo_matrix[Incomplete]: ... @overload # dtype: , format: "csc" -def identity(n: opt.AnyInt, dtype: npt.DTypeLike, format: _FmtCSC) -> csc_matrix[Incomplete]: ... +def identity(n: opt.AnyInt, dtype: _ToDType, format: _FmtCSC) -> csc_matrix[Incomplete]: ... @overload # dtype: , format: "csr" -def identity(n: opt.AnyInt, dtype: npt.DTypeLike, format: _FmtCSR) -> csr_matrix[Incomplete]: ... +def identity(n: opt.AnyInt, dtype: _ToDType, format: _FmtCSR) -> csr_matrix[Incomplete]: ... @overload # dtype: , format: "dok" -def identity(n: opt.AnyInt, dtype: npt.DTypeLike, format: _FmtDOK) -> dok_matrix[Incomplete]: ... +def identity(n: opt.AnyInt, dtype: _ToDType, format: _FmtDOK) -> dok_matrix[Incomplete]: ... @overload # dtype: , format: "lil" -def identity(n: opt.AnyInt, dtype: npt.DTypeLike, format: _FmtLIL) -> lil_matrix[Incomplete]: ... +def identity(n: opt.AnyInt, dtype: _ToDType, format: _FmtLIL) -> lil_matrix[Incomplete]: ... ### @overload # dtype: float64-like (default), format: "dia" | None @@ -1064,31 +1074,31 @@ def eye_array( # @overload # dtype: , format: "dia" | None def eye_array( - m: opt.AnyInt, n: opt.AnyInt | None = None, *, k: int = 0, dtype: npt.DTypeLike, format: _FmtDIA | None = None + m: opt.AnyInt, n: opt.AnyInt | None = None, *, k: int = 0, dtype: _ToDType, format: _FmtDIA | None = None ) -> _DIAArray[Incomplete]: ... @overload # dtype: , format: "bsr" def eye_array( - m: opt.AnyInt, n: opt.AnyInt | None = None, *, k: int = 0, dtype: npt.DTypeLike, format: _FmtBSR + m: opt.AnyInt, n: opt.AnyInt | None = None, *, k: int = 0, dtype: _ToDType, format: _FmtBSR ) -> _BSRArray[Incomplete]: ... @overload # dtype: , format: "coo" def eye_array( - m: opt.AnyInt, n: opt.AnyInt | None = None, *, k: int = 0, dtype: npt.DTypeLike, format: _FmtCOO + m: opt.AnyInt, n: opt.AnyInt | None = None, *, k: int = 0, dtype: _ToDType, format: _FmtCOO ) -> _COOArray2D[Incomplete]: ... @overload # dtype: , format: "csc" def eye_array( - m: opt.AnyInt, n: opt.AnyInt | None = None, *, k: int = 0, dtype: npt.DTypeLike, format: _FmtCSC + m: opt.AnyInt, n: opt.AnyInt | None = None, *, k: int = 0, dtype: _ToDType, format: _FmtCSC ) -> _CSCArray[Incomplete]: ... @overload # dtype: , format: "csr" def eye_array( - m: opt.AnyInt, n: opt.AnyInt | None = None, *, k: int = 0, dtype: npt.DTypeLike, format: _FmtCSR + m: opt.AnyInt, n: opt.AnyInt | None = None, *, k: int = 0, dtype: _ToDType, format: _FmtCSR ) -> _CSRArray2D[Incomplete]: ... @overload # dtype: , format: "dok" def eye_array( - m: opt.AnyInt, n: opt.AnyInt | None = None, *, k: int = 0, dtype: npt.DTypeLike, format: _FmtDOK + m: opt.AnyInt, n: opt.AnyInt | None = None, *, k: int = 0, dtype: _ToDType, format: _FmtDOK ) -> _DOKArray2D[Incomplete]: ... @overload # dtype: , format: "lil" def eye_array( - m: opt.AnyInt, n: opt.AnyInt | None = None, *, k: int = 0, dtype: npt.DTypeLike, format: _FmtLIL + m: opt.AnyInt, n: opt.AnyInt | None = None, *, k: int = 0, dtype: _ToDType, format: _FmtLIL ) -> _LILArray[Incomplete]: ... ### @@ -1245,31 +1255,31 @@ def eye( # @overload # dtype: , format: "dia" | None def eye( - m: opt.AnyInt, n: opt.AnyInt | None = None, k: int = 0, *, dtype: npt.DTypeLike, format: _FmtDIA | None = None + m: opt.AnyInt, n: opt.AnyInt | None = None, k: int = 0, *, dtype: _ToDType, format: _FmtDIA | None = None ) -> dia_matrix[Incomplete]: ... @overload # dtype: , format: "bsr" def eye( - m: opt.AnyInt, n: opt.AnyInt | None = None, k: int = 0, *, dtype: npt.DTypeLike, format: _FmtBSR + m: opt.AnyInt, n: opt.AnyInt | None = None, k: int = 0, *, dtype: _ToDType, format: _FmtBSR ) -> bsr_matrix[Incomplete]: ... @overload # dtype: , format: "coo" def eye( - m: opt.AnyInt, n: opt.AnyInt | None = None, k: int = 0, *, dtype: npt.DTypeLike, format: _FmtCOO + m: opt.AnyInt, n: opt.AnyInt | None = None, k: int = 0, *, dtype: _ToDType, format: _FmtCOO ) -> coo_matrix[Incomplete]: ... @overload # dtype: , format: "csc" def eye( - m: opt.AnyInt, n: opt.AnyInt | None = None, k: int = 0, *, dtype: npt.DTypeLike, format: _FmtCSC + m: opt.AnyInt, n: opt.AnyInt | None = None, k: int = 0, *, dtype: _ToDType, format: _FmtCSC ) -> csc_matrix[Incomplete]: ... @overload # dtype: , format: "csr" def eye( - m: opt.AnyInt, n: opt.AnyInt | None = None, k: int = 0, *, dtype: npt.DTypeLike, format: _FmtCSR + m: opt.AnyInt, n: opt.AnyInt | None = None, k: int = 0, *, dtype: _ToDType, format: _FmtCSR ) -> csr_matrix[Incomplete]: ... @overload # dtype: , format: "dok" def eye( - m: opt.AnyInt, n: opt.AnyInt | None = None, k: int = 0, *, dtype: npt.DTypeLike, format: _FmtDOK + m: opt.AnyInt, n: opt.AnyInt | None = None, k: int = 0, *, dtype: _ToDType, format: _FmtDOK ) -> dok_matrix[Incomplete]: ... @overload # dtype: , format: "lil" def eye( - m: opt.AnyInt, n: opt.AnyInt | None = None, k: int = 0, *, dtype: npt.DTypeLike, format: _FmtLIL + m: opt.AnyInt, n: opt.AnyInt | None = None, k: int = 0, *, dtype: _ToDType, format: _FmtLIL ) -> lil_matrix[Incomplete]: ... ### @@ -1378,19 +1388,19 @@ def kronsum(A: onp.ToComplex2D, B: onp.ToComplex2D, format: _Format | None = Non @overload # sparray, format: , dtype: def hstack(blocks: Seq[_CanStack[_T]], format: None = None, dtype: None = None) -> _T: ... @overload # sparray, format: "bsr", dtype: -def hstack(blocks: Seq[sparray[_SCT0]], format: _FmtBSR, dtype: None = None) -> _BSRArray[_SCT0]: ... +def hstack(blocks: Seq[_SpArray2D[_SCT]], format: _FmtBSR, dtype: None = None) -> _BSRArray[_SCT]: ... @overload # sparray, format: "coo", dtype: -def hstack(blocks: Seq[sparray[_SCT0]], format: _FmtCOO, dtype: None = None) -> _COOArray2D[_SCT0]: ... +def hstack(blocks: Seq[_SpArray2D[_SCT]], format: _FmtCOO, dtype: None = None) -> _COOArray2D[_SCT]: ... @overload # sparray, format: "csc", dtype: -def hstack(blocks: Seq[sparray[_SCT0]], format: _FmtCSC, dtype: None = None) -> _CSCArray[_SCT0]: ... +def hstack(blocks: Seq[_SpArray2D[_SCT]], format: _FmtCSC, dtype: None = None) -> _CSCArray[_SCT]: ... @overload # sparray, format: "csr", dtype: -def hstack(blocks: Seq[sparray[_SCT0]], format: _FmtCSR, dtype: None = None) -> _CSRArray2D[_SCT0]: ... +def hstack(blocks: Seq[_SpArray2D[_SCT]], format: _FmtCSR, dtype: None = None) -> _CSRArray2D[_SCT]: ... @overload # sparray, format: "dia", dtype: -def hstack(blocks: Seq[sparray[_SCT0]], format: _FmtDIA, dtype: None = None) -> _DIAArray[_SCT0]: ... +def hstack(blocks: Seq[_SpArray2D[_SCT]], format: _FmtDIA, dtype: None = None) -> _DIAArray[_SCT]: ... @overload # sparray, format: "dok", dtype: -def hstack(blocks: Seq[sparray[_SCT0]], format: _FmtDOK, dtype: None = None) -> _DOKArray2D[_SCT0]: ... +def hstack(blocks: Seq[_SpArray2D[_SCT]], format: _FmtDOK, dtype: None = None) -> _DOKArray2D[_SCT]: ... @overload # sparray, format: "lil", dtype: -def hstack(blocks: Seq[sparray[_SCT0]], format: _FmtLIL, dtype: None = None) -> _LILArray[_SCT0]: ... +def hstack(blocks: Seq[_SpArray2D[_SCT]], format: _FmtLIL, dtype: None = None) -> _LILArray[_SCT]: ... # @overload # sparray, format: , dtype: bool-like @@ -1484,44 +1494,44 @@ def hstack(blocks: Seq[sparray], format: _FmtLIL, dtype: onp.ToDType[_SCT0]) -> # @overload # sparray, format: , dtype: -def hstack(blocks: Seq[_CanStackAs[Any, _T]], format: None = None, *, dtype: npt.DTypeLike) -> _T: ... +def hstack(blocks: Seq[_CanStackAs[Any, _T]], format: None = None, *, dtype: _ToDType) -> _T: ... @overload # sparray, format: "bsr", dtype: -def hstack(blocks: Seq[sparray], format: _FmtBSR, dtype: npt.DTypeLike) -> _BSRArray: ... +def hstack(blocks: Seq[sparray], format: _FmtBSR, dtype: _ToDType) -> _BSRArray: ... @overload # sparray, format: "coo", dtype: -def hstack(blocks: Seq[sparray], format: _FmtCOO, dtype: npt.DTypeLike) -> _COOArray2D: ... +def hstack(blocks: Seq[sparray], format: _FmtCOO, dtype: _ToDType) -> _COOArray2D: ... @overload # sparray, format: "csc", dtype: -def hstack(blocks: Seq[sparray], format: _FmtCSC, dtype: npt.DTypeLike) -> _CSCArray: ... +def hstack(blocks: Seq[sparray], format: _FmtCSC, dtype: _ToDType) -> _CSCArray: ... @overload # sparray, format: "csr", dtype: -def hstack(blocks: Seq[sparray], format: _FmtCSR, dtype: npt.DTypeLike) -> _CSRArray2D: ... +def hstack(blocks: Seq[sparray], format: _FmtCSR, dtype: _ToDType) -> _CSRArray2D: ... @overload # sparray, format: "dia", dtype: -def hstack(blocks: Seq[sparray], format: _FmtDIA, dtype: npt.DTypeLike) -> _DIAArray: ... +def hstack(blocks: Seq[sparray], format: _FmtDIA, dtype: _ToDType) -> _DIAArray: ... @overload # sparray, format: "dok", dtype: -def hstack(blocks: Seq[sparray], format: _FmtDOK, dtype: npt.DTypeLike) -> _DOKArray2D: ... +def hstack(blocks: Seq[sparray], format: _FmtDOK, dtype: _ToDType) -> _DOKArray2D: ... @overload # sparray, format: "lil", dtype: -def hstack(blocks: Seq[sparray], format: _FmtLIL, dtype: npt.DTypeLike) -> _LILArray: ... +def hstack(blocks: Seq[sparray], format: _FmtLIL, dtype: _ToDType) -> _LILArray: ... # @overload -def hstack(blocks: Seq[_spbase], format: _Format, dtype: npt.DTypeLike | None = None) -> Incomplete: ... +def hstack(blocks: Seq[_spbase], format: _Format, dtype: _ToDType | None = None) -> Incomplete: ... ### # NOTE: keep in sync with `hstack` @overload # sparray, format: , dtype: def vstack(blocks: Seq[_CanStack[_T]], format: None = None, dtype: None = None) -> _T: ... @overload # sparray, format: "bsr", dtype: -def vstack(blocks: Seq[sparray[_SCT0]], format: _FmtBSR, dtype: None = None) -> _BSRArray[_SCT0]: ... +def vstack(blocks: Seq[_SpArray2D[_SCT]], format: _FmtBSR, dtype: None = None) -> _BSRArray[_SCT]: ... @overload # sparray, format: "coo", dtype: -def vstack(blocks: Seq[sparray[_SCT0]], format: _FmtCOO, dtype: None = None) -> _COOArray2D[_SCT0]: ... +def vstack(blocks: Seq[_SpArray2D[_SCT]], format: _FmtCOO, dtype: None = None) -> _COOArray2D[_SCT]: ... @overload # sparray, format: "csc", dtype: -def vstack(blocks: Seq[sparray[_SCT0]], format: _FmtCSC, dtype: None = None) -> _CSCArray[_SCT0]: ... +def vstack(blocks: Seq[_SpArray2D[_SCT]], format: _FmtCSC, dtype: None = None) -> _CSCArray[_SCT]: ... @overload # sparray, format: "csr", dtype: -def vstack(blocks: Seq[sparray[_SCT0]], format: _FmtCSR, dtype: None = None) -> _CSRArray2D[_SCT0]: ... +def vstack(blocks: Seq[_SpArray2D[_SCT]], format: _FmtCSR, dtype: None = None) -> _CSRArray2D[_SCT]: ... @overload # sparray, format: "dia", dtype: -def vstack(blocks: Seq[sparray[_SCT0]], format: _FmtDIA, dtype: None = None) -> _DIAArray[_SCT0]: ... +def vstack(blocks: Seq[_SpArray2D[_SCT]], format: _FmtDIA, dtype: None = None) -> _DIAArray[_SCT]: ... @overload # sparray, format: "dok", dtype: -def vstack(blocks: Seq[sparray[_SCT0]], format: _FmtDOK, dtype: None = None) -> _DOKArray2D[_SCT0]: ... +def vstack(blocks: Seq[_SpArray2D[_SCT]], format: _FmtDOK, dtype: None = None) -> _DOKArray2D[_SCT]: ... @overload # sparray, format: "lil", dtype: -def vstack(blocks: Seq[sparray[_SCT0]], format: _FmtLIL, dtype: None = None) -> _LILArray[_SCT0]: ... +def vstack(blocks: Seq[_SpArray2D[_SCT]], format: _FmtLIL, dtype: None = None) -> _LILArray[_SCT]: ... # @overload # sparray, format: , dtype: bool-like @@ -1615,138 +1625,327 @@ def vstack(blocks: Seq[sparray], format: _FmtLIL, dtype: onp.ToDType[_SCT0]) -> # @overload # sparray, format: , dtype: -def vstack(blocks: Seq[_CanStackAs[Any, _T]], format: None = None, *, dtype: npt.DTypeLike) -> _T: ... +def vstack(blocks: Seq[_CanStackAs[Any, _T]], format: None = None, *, dtype: _ToDType) -> _T: ... @overload # sparray, format: "bsr", dtype: -def vstack(blocks: Seq[sparray], format: _FmtBSR, dtype: npt.DTypeLike) -> _BSRArray: ... +def vstack(blocks: Seq[sparray], format: _FmtBSR, dtype: _ToDType) -> _BSRArray: ... @overload # sparray, format: "coo", dtype: -def vstack(blocks: Seq[sparray], format: _FmtCOO, dtype: npt.DTypeLike) -> _COOArray2D: ... +def vstack(blocks: Seq[sparray], format: _FmtCOO, dtype: _ToDType) -> _COOArray2D: ... @overload # sparray, format: "csc", dtype: -def vstack(blocks: Seq[sparray], format: _FmtCSC, dtype: npt.DTypeLike) -> _CSCArray: ... +def vstack(blocks: Seq[sparray], format: _FmtCSC, dtype: _ToDType) -> _CSCArray: ... @overload # sparray, format: "csr", dtype: -def vstack(blocks: Seq[sparray], format: _FmtCSR, dtype: npt.DTypeLike) -> _CSRArray2D: ... +def vstack(blocks: Seq[sparray], format: _FmtCSR, dtype: _ToDType) -> _CSRArray2D: ... @overload # sparray, format: "dia", dtype: -def vstack(blocks: Seq[sparray], format: _FmtDIA, dtype: npt.DTypeLike) -> _DIAArray: ... +def vstack(blocks: Seq[sparray], format: _FmtDIA, dtype: _ToDType) -> _DIAArray: ... @overload # sparray, format: "dok", dtype: -def vstack(blocks: Seq[sparray], format: _FmtDOK, dtype: npt.DTypeLike) -> _DOKArray2D: ... +def vstack(blocks: Seq[sparray], format: _FmtDOK, dtype: _ToDType) -> _DOKArray2D: ... @overload # sparray, format: "lil", dtype: -def vstack(blocks: Seq[sparray], format: _FmtLIL, dtype: npt.DTypeLike) -> _LILArray: ... +def vstack(blocks: Seq[sparray], format: _FmtLIL, dtype: _ToDType) -> _LILArray: ... # @overload -def vstack(blocks: Seq[_spbase], format: _Format, dtype: npt.DTypeLike | None = None) -> Incomplete: ... +def vstack(blocks: Seq[_spbase], format: _Format, dtype: _ToDType | None = None) -> Incomplete: ... ### -# TODO(jorenham): Use `_CanStack` here, which requires a way to map matrix types to array types. -@overload # blocks: , format: , dtype: -def block_array(blocks: _ToBlocks[_SCT], *, format: _FmtCOO | None = None, dtype: None = None) -> _COOArray2D[_SCT]: ... -@overload # blocks: , format: , dtype: -def block_array(blocks: _ToBlocks[_SCT], *, format: _FmtNonCOO, dtype: None = None) -> _SpArray2D[_SCT]: ... - -# -@overload # blocks: , format: , dtype: -def block_array(blocks: _ToBlocks, *, format: _FmtCOO | None = None, dtype: onp.ToDType[_SCT]) -> _COOArray2D[_SCT]: ... -@overload # blocks: , format: , dtype: -def block_array(blocks: _ToBlocks, *, format: _FmtNonCOO, dtype: onp.ToDType[_SCT]) -> _SpArray2D[_SCT]: ... - -# -@overload # blocks: , format: , dtype: -def block_array(blocks: _ToBlocks, *, format: _FmtCOO | None = None, dtype: npt.DTypeLike) -> _COOArray2D: ... -@overload # blocks: , format: , dtype: -def block_array(blocks: _ToBlocks, *, format: _FmtNonCOO, dtype: npt.DTypeLike) -> _SpArray2D: ... +@overload # blocks: , format: , dtype: +def block_array(blocks: _ToBlocksCanStack[_T], *, format: None = None, dtype: None = None) -> _T: ... +@overload # blocks: , format: "bsr", dtype: +def block_array(blocks: _ToBlocksSPArray[_SCT], *, format: _FmtBSR, dtype: None = None) -> _BSRArray[_SCT]: ... +@overload # blocks: , format: "coo", dtype: +def block_array(blocks: _ToBlocksSPArray[_SCT], *, format: _FmtCOO, dtype: None = None) -> _COOArray2D[_SCT]: ... +@overload # blocks: , format: "csc", dtype: +def block_array(blocks: _ToBlocksSPArray[_SCT], *, format: _FmtCSC, dtype: None = None) -> _CSCArray[_SCT]: ... +@overload # blocks: , format: "csr", dtype: +def block_array(blocks: _ToBlocksSPArray[_SCT], *, format: _FmtCSR, dtype: None = None) -> _CSRArray2D[_SCT]: ... +@overload # blocks: , format: "dia", dtype: +def block_array(blocks: _ToBlocksSPArray[_SCT], *, format: _FmtDIA, dtype: None = None) -> _DIAArray[_SCT]: ... +@overload # blocks: , format: "dok", dtype: +def block_array(blocks: _ToBlocksSPArray[_SCT], *, format: _FmtDOK, dtype: None = None) -> _DOKArray2D[_SCT]: ... +@overload # blocks: , format: "lil", dtype: +def block_array(blocks: _ToBlocksSPArray[_SCT], *, format: _FmtLIL, dtype: None = None) -> _LILArray[_SCT]: ... + +# +@overload # blocks: , format: , dtype: bool-like +def block_array(blocks: _ToBlocksCanStackAs[np.bool_, _T], *, format: None = None, dtype: onp.AnyBoolDType) -> _T: ... +@overload # blocks: , format: "bsr", dtype: bool-like +def block_array(blocks: _ToBlocksUnkown, *, format: _FmtBSR, dtype: onp.AnyBoolDType) -> _BSRArray[np.bool_]: ... +@overload # blocks: , format: "coo", dtype: bool-like +def block_array(blocks: _ToBlocksUnkown, *, format: _FmtCOO, dtype: onp.AnyBoolDType) -> _COOArray2D[np.bool_]: ... +@overload # blocks: , format: "csc", dtype: bool-like +def block_array(blocks: _ToBlocksUnkown, *, format: _FmtCSC, dtype: onp.AnyBoolDType) -> _CSCArray[np.bool_]: ... +@overload # blocks: , format: "csr", dtype: bool-like +def block_array(blocks: _ToBlocksUnkown, *, format: _FmtCSR, dtype: onp.AnyBoolDType) -> _CSRArray2D[np.bool_]: ... +@overload # blocks: , format: "dia", dtype: bool-like +def block_array(blocks: _ToBlocksUnkown, *, format: _FmtDIA, dtype: onp.AnyBoolDType) -> _DIAArray[np.bool_]: ... +@overload # blocks: , format: "dok", dtype: bool-like +def block_array(blocks: _ToBlocksUnkown, *, format: _FmtDOK, dtype: onp.AnyBoolDType) -> _DOKArray2D[np.bool_]: ... +@overload # blocks: , format: "lil", dtype: bool-like +def block_array(blocks: _ToBlocksUnkown, *, format: _FmtLIL, dtype: onp.AnyBoolDType) -> _LILArray[np.bool_]: ... + +# +@overload # blocks: , format: , dtype: int-like +def block_array(blocks: _ToBlocksCanStackAs[np.int64, _T], *, format: None = None, dtype: onp.AnyIntDType) -> _T: ... +@overload # blocks: , format: "bsr", dtype: int-like +def block_array(blocks: _ToBlocksUnkown, *, format: _FmtBSR, dtype: onp.AnyIntDType) -> _BSRArray[np.int_]: ... +@overload # blocks: , format: "coo", dtype: int-like +def block_array(blocks: _ToBlocksUnkown, *, format: _FmtCOO, dtype: onp.AnyIntDType) -> _COOArray2D[np.int_]: ... +@overload # blocks: , format: "csc", dtype: int-like +def block_array(blocks: _ToBlocksUnkown, *, format: _FmtCSC, dtype: onp.AnyIntDType) -> _CSCArray[np.int_]: ... +@overload # blocks: , format: "csr", dtype: int-like +def block_array(blocks: _ToBlocksUnkown, *, format: _FmtCSR, dtype: onp.AnyIntDType) -> _CSRArray2D[np.int_]: ... +@overload # blocks: , format: "dia", dtype: int-like +def block_array(blocks: _ToBlocksUnkown, *, format: _FmtDIA, dtype: onp.AnyIntDType) -> _DIAArray[np.int_]: ... +@overload # blocks: , format: "dok", dtype: int-like +def block_array(blocks: _ToBlocksUnkown, *, format: _FmtDOK, dtype: onp.AnyIntDType) -> _DOKArray2D[np.int_]: ... +@overload # blocks: , format: "lil", dtype: int-like +def block_array(blocks: _ToBlocksUnkown, *, format: _FmtLIL, dtype: onp.AnyIntDType) -> _LILArray[np.int_]: ... + +# +@overload # blocks: , format: , dtype: float64-like +def block_array(blocks: _ToBlocksCanStackAs[np.float64, _T], *, format: None = None, dtype: onp.AnyFloat64DType) -> _T: ... +@overload # blocks: , format: "bsr", dtype: float64-like +def block_array(blocks: _ToBlocksUnkown, *, format: _FmtBSR, dtype: onp.AnyFloat64DType) -> _BSRArray[np.float64]: ... +@overload # blocks: , format: "coo", dtype: float64-like +def block_array(blocks: _ToBlocksUnkown, *, format: _FmtCOO, dtype: onp.AnyFloat64DType) -> _COOArray2D[np.float64]: ... +@overload # blocks: , format: "csc", dtype: float64-like +def block_array(blocks: _ToBlocksUnkown, *, format: _FmtCSC, dtype: onp.AnyFloat64DType) -> _CSCArray[np.float64]: ... +@overload # blocks: , format: "csr", dtype: float64-like +def block_array(blocks: _ToBlocksUnkown, *, format: _FmtCSR, dtype: onp.AnyFloat64DType) -> _CSRArray2D[np.float64]: ... +@overload # blocks: , format: "dia", dtype: float64-like +def block_array(blocks: _ToBlocksUnkown, *, format: _FmtDIA, dtype: onp.AnyFloat64DType) -> _DIAArray[np.float64]: ... +@overload # blocks: , format: "dok", dtype: float64-like +def block_array(blocks: _ToBlocksUnkown, *, format: _FmtDOK, dtype: onp.AnyFloat64DType) -> _DOKArray2D[np.float64]: ... +@overload # blocks: , format: "lil", dtype: float64-like +def block_array(blocks: _ToBlocksUnkown, *, format: _FmtLIL, dtype: onp.AnyFloat64DType) -> _LILArray[np.float64]: ... + +# +@overload # blocks: , format: , dtype: complex128-like +def block_array(blocks: _ToBlocksCanStackAs[np.complex128, _T], *, format: None = None, dtype: onp.AnyComplex128DType) -> _T: ... +@overload # blocks: , format: "bsr", dtype: complex128-like +def block_array(blocks: _ToBlocksUnkown, *, format: _FmtBSR, dtype: onp.AnyComplex128DType) -> _BSRArray[np.complex128]: ... +@overload # blocks: , format: "coo", dtype: complex128-like +def block_array(blocks: _ToBlocksUnkown, *, format: _FmtCOO, dtype: onp.AnyComplex128DType) -> _COOArray2D[np.complex128]: ... +@overload # blocks: , format: "csc", dtype: complex128-like +def block_array(blocks: _ToBlocksUnkown, *, format: _FmtCSC, dtype: onp.AnyComplex128DType) -> _CSCArray[np.complex128]: ... +@overload # blocks: , format: "csr", dtype: complex128-like +def block_array(blocks: _ToBlocksUnkown, *, format: _FmtCSR, dtype: onp.AnyComplex128DType) -> _CSRArray2D[np.complex128]: ... +@overload # blocks: , format: "dia", dtype: complex128-like +def block_array(blocks: _ToBlocksUnkown, *, format: _FmtDIA, dtype: onp.AnyComplex128DType) -> _DIAArray[np.complex128]: ... +@overload # blocks: , format: "dok", dtype: complex128-like +def block_array(blocks: _ToBlocksUnkown, *, format: _FmtDOK, dtype: onp.AnyComplex128DType) -> _DOKArray2D[np.complex128]: ... +@overload # blocks: , format: "lil", dtype: complex128-like +def block_array(blocks: _ToBlocksUnkown, *, format: _FmtLIL, dtype: onp.AnyComplex128DType) -> _LILArray[np.complex128]: ... + +# +@overload # blocks: , format: , dtype: +def block_array(blocks: _ToBlocksCanStackAs[_SCT0, _T], *, format: None = None, dtype: onp.ToDType[_SCT0]) -> _T: ... +@overload # blocks: , format: "bsr", dtype: +def block_array(blocks: _ToBlocksUnkown, *, format: _FmtBSR, dtype: onp.ToDType[_SCT0]) -> _BSRArray[_SCT0]: ... +@overload # blocks: , format: "coo", dtype: +def block_array(blocks: _ToBlocksUnkown, *, format: _FmtCOO, dtype: onp.ToDType[_SCT0]) -> _COOArray2D[_SCT0]: ... +@overload # blocks: , format: "csc", dtype: +def block_array(blocks: _ToBlocksUnkown, *, format: _FmtCSC, dtype: onp.ToDType[_SCT0]) -> _CSCArray[_SCT0]: ... +@overload # blocks: , format: "csr", dtype: +def block_array(blocks: _ToBlocksUnkown, *, format: _FmtCSR, dtype: onp.ToDType[_SCT0]) -> _CSRArray2D[_SCT0]: ... +@overload # blocks: , format: "dia", dtype: +def block_array(blocks: _ToBlocksUnkown, *, format: _FmtDIA, dtype: onp.ToDType[_SCT0]) -> _DIAArray[_SCT0]: ... +@overload # blocks: , format: "dok", dtype: +def block_array(blocks: _ToBlocksUnkown, *, format: _FmtDOK, dtype: onp.ToDType[_SCT0]) -> _DOKArray2D[_SCT0]: ... +@overload # blocks: , format: "lil", dtype: +def block_array(blocks: _ToBlocksUnkown, *, format: _FmtLIL, dtype: onp.ToDType[_SCT0]) -> _LILArray[_SCT0]: ... + +# +@overload # blocks: , format: , dtype: +def block_array(blocks: _ToBlocksCanStackAs[Any, _T], *, format: None = None, dtype: _ToDType | None = None) -> _T: ... +@overload # blocks: , format: "bsr", dtype: +def block_array(blocks: _ToBlocksUnkown, *, format: _FmtBSR, dtype: _ToDType | None = None) -> _BSRArray: ... +@overload # blocks: , format: "coo", dtype: +def block_array(blocks: _ToBlocksUnkown, *, format: _FmtCOO, dtype: _ToDType | None = None) -> _COOArray2D: ... +@overload # blocks: , format: "csc", dtype: +def block_array(blocks: _ToBlocksUnkown, *, format: _FmtCSC, dtype: _ToDType | None = None) -> _CSCArray: ... +@overload # blocks: , format: "csr", dtype: +def block_array(blocks: _ToBlocksUnkown, *, format: _FmtCSR, dtype: _ToDType | None = None) -> _CSRArray2D: ... +@overload # blocks: , format: "dia", dtype: +def block_array(blocks: _ToBlocksUnkown, *, format: _FmtDIA, dtype: _ToDType | None = None) -> _DIAArray: ... +@overload # blocks: , format: "dok", dtype: +def block_array(blocks: _ToBlocksUnkown, *, format: _FmtDOK, dtype: _ToDType | None = None) -> _DOKArray2D: ... +@overload # blocks: , format: "lil", dtype: +def block_array(blocks: _ToBlocksUnkown, *, format: _FmtLIL, dtype: _ToDType | None = None) -> _LILArray: ... ### -# TODO(jorenham): Use `_CanStack` here, which requires a way to map array types to matrix types. -@overload # blocks: , format: , dtype: -def bmat(blocks: Seq[Seq[sparray[_SCT]]], format: _FmtCOO | None = None, dtype: None = None) -> _COOArray2D[_SCT]: ... -@overload # blocks: , format: , dtype: -def bmat(blocks: Seq[Seq[sparray[_SCT]]], format: _FmtNonCOO, dtype: None = None) -> _SpArray2D[_SCT]: ... - -# -@overload # blocks: , format: , dtype: -def bmat(blocks: Seq[Seq[spmatrix[_SCT]]], format: _FmtCOO | None = None, dtype: None = None) -> coo_matrix[_SCT]: ... +@overload # blocks: , format: , dtype: +def bmat(blocks: _ToBlocksCanStack[_T], format: None = None, dtype: None = None) -> _T: ... @overload # blocks: , format: , dtype: -def bmat(blocks: Seq[Seq[spmatrix[_SCT]]], format: _FmtNonCOO, dtype: None = None) -> spmatrix[_SCT]: ... +def bmat(blocks: _ToBlocksSPMatrix[_SCT], format: _Format, dtype: None = None) -> spmatrix[_SCT]: ... # -@overload # blocks: , format: , dtype: -def bmat( - blocks: _ToBlocks, format: _FmtCOO | None = None, *, dtype: onp.ToDType[_SCT] -) -> _COOArray2D[_SCT] | coo_matrix[_SCT]: ... +@overload # blocks: , format: , dtype: +def bmat(blocks: _ToBlocksCanStackAs[_SCT0, _T], format: None = None, *, dtype: onp.ToDType[_SCT0]) -> _T: ... +@overload # blocks: , format: , dtype: +def bmat(blocks: _ToBlocksSPMatrix[_SCT0], format: _Format, dtype: onp.ToDType[_SCT0]) -> spmatrix[_SCT0]: ... @overload # blocks: , format: , dtype: -def bmat(blocks: _ToBlocks, format: _FmtNonCOO, dtype: onp.ToDType[_SCT]) -> _SpBase2D[_SCT]: ... +def bmat(blocks: _ToBlocksUnkown, format: _Format, dtype: onp.ToDType[_SCT0]) -> spmatrix[_SCT0] | _SpArray2D[_SCT0]: ... # -@overload # blocks: , format: , dtype: -def bmat( - blocks: _ToBlocks[_SCT], format: _FmtCOO | None = None, dtype: npt.DTypeLike | None = None -) -> _COOArray2D[_SCT] | coo_matrix[_SCT]: ... +@overload # blocks: , format: , dtype: +def bmat(blocks: _ToBlocksCanStackAs[Any, _T], format: None = None, *, dtype: _ToDType) -> _T: ... +@overload # blocks: , format: , dtype: +def bmat(blocks: _ToBlocksSPMatrix, format: _Format, *, dtype: _ToDType) -> spmatrix: ... @overload # blocks: , format: , dtype: -def bmat(blocks: _ToBlocks[_SCT], format: _FmtNonCOO, dtype: npt.DTypeLike) -> _SpBase2D[_SCT]: ... +def bmat(blocks: _ToBlocksUnkown, format: _Format, *, dtype: _ToDType) -> spmatrix | _SpArray2D: ... ### -# TODO(jorenham): Add support for non-COO formats. @overload # mats: , format: , dtype: None def block_diag(mats: Iterable[sparray[_SCT]], format: _FmtCOO | None = None, dtype: None = None) -> _COOArray2D[_SCT]: ... - -# -@overload # mats: , format: , dtype: None -def block_diag(mats: Iterable[spmatrix[_SCT]], format: _FmtCOO | None = None, dtype: None = None) -> coo_matrix[_SCT]: ... - -# -@overload # mats: , format: , dtype: None +@overload # mats: , format: "bsr", dtype: None +def block_diag(mats: Iterable[sparray[_SCT]], format: _FmtBSR, dtype: None = None) -> _BSRArray[_SCT]: ... +@overload # mats: , format: "csc", dtype: None +def block_diag(mats: Iterable[sparray[_SCT]], format: _FmtCSC, dtype: None = None) -> _CSCArray[_SCT]: ... +@overload # mats: , format: "csr", dtype: None +def block_diag(mats: Iterable[sparray[_SCT]], format: _FmtCSR, dtype: None = None) -> _CSRArray2D[_SCT]: ... +@overload # mats: , format: "dia", dtype: None +def block_diag(mats: Iterable[sparray[_SCT]], format: _FmtDIA, dtype: None = None) -> _DIAArray[_SCT]: ... +@overload # mats: , format: "dok", dtype: None +def block_diag(mats: Iterable[sparray[_SCT]], format: _FmtDOK, dtype: None = None) -> _DOKArray2D[_SCT]: ... +@overload # mats: , format: "lil", dtype: None +def block_diag(mats: Iterable[sparray[_SCT]], format: _FmtLIL, dtype: None = None) -> _LILArray[_SCT]: ... + +# +@overload # mats: , format: , dtype: bool-like +def block_diag(mats: Iterable[sparray], format: _FmtCOO | None = None, *, dtype: onp.AnyBoolDType) -> _COOArray2D[np.bool_]: ... +@overload # mats: , format: "bsr", dtype: bool-like +def block_diag(mats: Iterable[sparray], format: _FmtBSR, dtype: onp.AnyBoolDType) -> _BSRArray[np.bool_]: ... +@overload # mats: , format: "csc", dtype: bool-like +def block_diag(mats: Iterable[sparray], format: _FmtCSC, dtype: onp.AnyBoolDType) -> _CSCArray[np.bool_]: ... +@overload # mats: , format: "csr", dtype: bool-like +def block_diag(mats: Iterable[sparray], format: _FmtCSR, dtype: onp.AnyBoolDType) -> _CSRArray2D[np.bool_]: ... +@overload # mats: , format: "dia", dtype: bool-like +def block_diag(mats: Iterable[sparray], format: _FmtDIA, dtype: onp.AnyBoolDType) -> _DIAArray[np.bool_]: ... +@overload # mats: , format: "dok", dtype: bool-like +def block_diag(mats: Iterable[sparray], format: _FmtDOK, dtype: onp.AnyBoolDType) -> _DOKArray2D[np.bool_]: ... +@overload # mats: , format: "lil", dtype: bool-like +def block_diag(mats: Iterable[sparray], format: _FmtLIL, dtype: onp.AnyBoolDType) -> _LILArray[np.bool_]: ... + +# +@overload # mats: , format: , dtype: int-like +def block_diag(mats: Iterable[sparray], format: _FmtCOO | None = None, *, dtype: onp.AnyIntDType) -> _COOArray2D[np.int_]: ... +@overload # mats: , format: "bsr", dtype: int-like +def block_diag(mats: Iterable[sparray], format: _FmtBSR, dtype: onp.AnyIntDType) -> _BSRArray[np.int_]: ... +@overload # mats: , format: "csc", dtype: int-like +def block_diag(mats: Iterable[sparray], format: _FmtCSC, dtype: onp.AnyIntDType) -> _CSCArray[np.int_]: ... +@overload # mats: , format: "csr", dtype: int-like +def block_diag(mats: Iterable[sparray], format: _FmtCSR, dtype: onp.AnyIntDType) -> _CSRArray2D[np.int_]: ... +@overload # mats: , format: "dia", dtype: int-like +def block_diag(mats: Iterable[sparray], format: _FmtDIA, dtype: onp.AnyIntDType) -> _DIAArray[np.int_]: ... +@overload # mats: , format: "dok", dtype: int-like +def block_diag(mats: Iterable[sparray], format: _FmtDOK, dtype: onp.AnyIntDType) -> _DOKArray2D[np.int_]: ... +@overload # mats: , format: "lil", dtype: int-like +def block_diag(mats: Iterable[sparray], format: _FmtLIL, dtype: onp.AnyIntDType) -> _LILArray[np.int_]: ... + +# +@overload # mats: , format: , dtype: float64-like def block_diag( - mats: Iterable[_spbase[_SCT] | onp.ArrayND[_SCT]], format: _FmtCOO | None = None, dtype: None = None -) -> _COOArray2D[_SCT] | coo_matrix[_SCT]: ... - -# -@overload # mats: , format: , dtype: (keyword) -def block_diag(mats: Iterable[sparray], format: _FmtCOO | None = None, *, dtype: onp.ToDType[_SCT]) -> _COOArray2D[_SCT]: ... -@overload # mats: , format: , dtype: (positional) -def block_diag(mats: Iterable[sparray], format: _FmtCOO | None, dtype: onp.ToDType[_SCT]) -> _COOArray2D[_SCT]: ... - -# -@overload # mats: , format: , dtype: (keyword) -def block_diag( - mats: Iterable[spmatrix | onp.ArrayND[_Numeric] | complex | Seq[onp.ToComplex] | Seq[onp.ToComplex1D]], - format: _FmtCOO | None = None, - *, - dtype: onp.ToDType[_SCT], -) -> coo_matrix[_SCT]: ... -@overload # mats: , format: , dtype: (positional) + mats: Iterable[sparray], format: _FmtCOO | None = None, *, dtype: onp.AnyFloat64DType +) -> _COOArray2D[np.float64]: ... +@overload # mats: , format: "bsr", dtype: float64-like +def block_diag(mats: Iterable[sparray], format: _FmtBSR, dtype: onp.AnyFloat64DType) -> _BSRArray[np.float64]: ... +@overload # mats: , format: "csc", dtype: float64-like +def block_diag(mats: Iterable[sparray], format: _FmtCSC, dtype: onp.AnyFloat64DType) -> _CSCArray[np.float64]: ... +@overload # mats: , format: "csr", dtype: float64-like +def block_diag(mats: Iterable[sparray], format: _FmtCSR, dtype: onp.AnyFloat64DType) -> _CSRArray2D[np.float64]: ... +@overload # mats: , format: "dia", dtype: float64-like +def block_diag(mats: Iterable[sparray], format: _FmtDIA, dtype: onp.AnyFloat64DType) -> _DIAArray[np.float64]: ... +@overload # mats: , format: "dok", dtype: float64-like +def block_diag(mats: Iterable[sparray], format: _FmtDOK, dtype: onp.AnyFloat64DType) -> _DOKArray2D[np.float64]: ... +@overload # mats: , format: "lil", dtype: float64-like +def block_diag(mats: Iterable[sparray], format: _FmtLIL, dtype: onp.AnyFloat64DType) -> _LILArray[np.float64]: ... + +# +@overload # mats: , format: , dtype: complex128-like def block_diag( - mats: Iterable[spmatrix | onp.ArrayND[_Numeric] | complex | Seq[onp.ToComplex] | Seq[onp.ToComplex1D]], - format: _FmtCOO | None, - dtype: onp.ToDType[_SCT], -) -> coo_matrix[_SCT]: ... + mats: Iterable[sparray], format: _FmtCOO | None = None, *, dtype: onp.AnyComplex128DType +) -> _COOArray2D[np.complex128]: ... +@overload # mats: , format: "bsr", dtype: complex128-like +def block_diag(mats: Iterable[sparray], format: _FmtBSR, dtype: onp.AnyComplex128DType) -> _BSRArray[np.complex128]: ... +@overload # mats: , format: "csc", dtype: complex128-like +def block_diag(mats: Iterable[sparray], format: _FmtCSC, dtype: onp.AnyComplex128DType) -> _CSCArray[np.complex128]: ... +@overload # mats: , format: "csr", dtype: complex128-like +def block_diag(mats: Iterable[sparray], format: _FmtCSR, dtype: onp.AnyComplex128DType) -> _CSRArray2D[np.complex128]: ... +@overload # mats: , format: "dia", dtype: complex128-like +def block_diag(mats: Iterable[sparray], format: _FmtDIA, dtype: onp.AnyComplex128DType) -> _DIAArray[np.complex128]: ... +@overload # mats: , format: "dok", dtype: complex128-like +def block_diag(mats: Iterable[sparray], format: _FmtDOK, dtype: onp.AnyComplex128DType) -> _DOKArray2D[np.complex128]: ... +@overload # mats: , format: "lil", dtype: complex128-like +def block_diag(mats: Iterable[sparray], format: _FmtLIL, dtype: onp.AnyComplex128DType) -> _LILArray[np.complex128]: ... + +# +@overload # mats: , format: , dtype: +def block_diag(mats: Iterable[sparray], format: _FmtCOO | None = None, *, dtype: onp.ToDType[_SCT]) -> _COOArray2D[_SCT]: ... +@overload # mats: , format: "bsr", dtype: +def block_diag(mats: Iterable[sparray], format: _FmtBSR, dtype: onp.ToDType[_SCT]) -> _BSRArray[_SCT]: ... +@overload # mats: , format: "csc", dtype: +def block_diag(mats: Iterable[sparray], format: _FmtCSC, dtype: onp.ToDType[_SCT]) -> _CSCArray[_SCT]: ... +@overload # mats: , format: "csr", dtype: +def block_diag(mats: Iterable[sparray], format: _FmtCSR, dtype: onp.ToDType[_SCT]) -> _CSRArray2D[_SCT]: ... +@overload # mats: , format: "dia", dtype: +def block_diag(mats: Iterable[sparray], format: _FmtDIA, dtype: onp.ToDType[_SCT]) -> _DIAArray[_SCT]: ... +@overload # mats: , format: "dok", dtype: +def block_diag(mats: Iterable[sparray], format: _FmtDOK, dtype: onp.ToDType[_SCT]) -> _DOKArray2D[_SCT]: ... +@overload # mats: , format: "lil", dtype: +def block_diag(mats: Iterable[sparray], format: _FmtLIL, dtype: onp.ToDType[_SCT]) -> _LILArray[_SCT]: ... # -@overload # mats: , format: , dtype: (keyword) +@overload # mats: , format: , dtype: None def block_diag( - mats: Iterable[_spbase | onp.ArrayND[_Numeric] | complex | Seq[onp.ToComplex] | Seq[onp.ToComplex1D]], - format: _FmtCOO | None = None, - *, - dtype: onp.ToDType[_SCT], + mats: _ToMatsDiag[_SCT], format: _FmtCOO | None = None, dtype: None = None ) -> _COOArray2D[_SCT] | coo_matrix[_SCT]: ... -@overload # mats: , format: , dtype: (positional) +@overload # mats: , format: "bsr", dtype: None +def block_diag(mats: _ToMatsDiag[_SCT], format: _FmtBSR, dtype: None = None) -> _BSRArray[_SCT] | bsr_matrix[_SCT]: ... +@overload # mats: , format: "csc", dtype: None +def block_diag(mats: _ToMatsDiag[_SCT], format: _FmtCSC, dtype: None = None) -> _CSCArray[_SCT] | csc_matrix[_SCT]: ... +@overload # mats: , format: "csr", dtype: None +def block_diag(mats: _ToMatsDiag[_SCT], format: _FmtCSR, dtype: None = None) -> _CSRArray2D[_SCT] | csr_matrix[_SCT]: ... +@overload # mats: , format: "dia", dtype: None +def block_diag(mats: _ToMatsDiag[_SCT], format: _FmtDIA, dtype: None = None) -> _DIAArray[_SCT] | dia_matrix[_SCT]: ... +@overload # mats: , format: "dok", dtype: None +def block_diag(mats: _ToMatsDiag[_SCT], format: _FmtDOK, dtype: None = None) -> _DOKArray2D[_SCT] | dok_matrix[_SCT]: ... +@overload # mats: , format: "lil", dtype: None +def block_diag(mats: _ToMatsDiag[_SCT], format: _FmtLIL, dtype: None = None) -> _LILArray[_SCT] | lil_matrix[_SCT]: ... + +# +@overload # mats: , format: , dtype: def block_diag( - mats: Iterable[_spbase | onp.ArrayND[_Numeric] | complex | Seq[onp.ToComplex] | Seq[onp.ToComplex1D]], - format: _FmtCOO | None, - dtype: onp.ToDType[_SCT], + mats: _ToMatsDiagUnknown, format: _FmtCOO | None = None, *, dtype: onp.ToDType[_SCT] ) -> _COOArray2D[_SCT] | coo_matrix[_SCT]: ... +@overload # mats: , format: "bsr", dtype: +def block_diag(mats: _ToMatsDiagUnknown, format: _FmtBSR, dtype: onp.ToDType[_SCT]) -> _BSRArray[_SCT] | bsr_matrix[_SCT]: ... +@overload # mats: , format: "csc", dtype: +def block_diag(mats: _ToMatsDiagUnknown, format: _FmtCSC, dtype: onp.ToDType[_SCT]) -> _CSCArray[_SCT] | csc_matrix[_SCT]: ... +@overload # mats: , format: "csr", dtype: +def block_diag(mats: _ToMatsDiagUnknown, format: _FmtCSR, dtype: onp.ToDType[_SCT]) -> _CSRArray2D[_SCT] | csr_matrix[_SCT]: ... +@overload # mats: , format: "dia", dtype: +def block_diag(mats: _ToMatsDiagUnknown, format: _FmtDIA, dtype: onp.ToDType[_SCT]) -> _DIAArray[_SCT] | dia_matrix[_SCT]: ... +@overload # mats: , format: "dok", dtype: +def block_diag(mats: _ToMatsDiagUnknown, format: _FmtDOK, dtype: onp.ToDType[_SCT]) -> _DOKArray2D[_SCT] | dok_matrix[_SCT]: ... +@overload # mats: , format: "lil", dtype: +def block_diag(mats: _ToMatsDiagUnknown, format: _FmtLIL, dtype: onp.ToDType[_SCT]) -> _LILArray[_SCT] | lil_matrix[_SCT]: ... # @overload # mats: , format: , dtype: def block_diag( - mats: Iterable[_spbase | onp.ArrayND[_Numeric] | complex | Seq[onp.ToComplex] | Seq[onp.ToComplex1D]], - format: _FmtCOO | None = None, - dtype: npt.DTypeLike | None = None, -) -> _COOArray2D[_SCT] | coo_matrix[Any]: ... -@overload # mats: , format: , dtype: -def block_diag( - mats: Iterable[_spbase | onp.ArrayND[_Numeric] | complex | Seq[onp.ToComplex] | Seq[onp.ToComplex1D]], - format: _FmtNonCOO | None = None, - dtype: npt.DTypeLike | None = None, -) -> Incomplete: ... + mats: _ToMatsDiagUnknown, format: _FmtCOO | None = None, dtype: _ToDType | None = None +) -> _COOArray2D | coo_matrix: ... +@overload # mats: , format: "bsr", dtype: +def block_diag(mats: _ToMatsDiagUnknown, format: _FmtBSR, dtype: _ToDType | None = None) -> _BSRArray | bsr_matrix: ... +@overload # mats: , format: "csc", dtype: +def block_diag(mats: _ToMatsDiagUnknown, format: _FmtCSC, dtype: _ToDType | None = None) -> _CSCArray | csc_matrix: ... +@overload # mats: , format: "csr", dtype: +def block_diag(mats: _ToMatsDiagUnknown, format: _FmtCSR, dtype: _ToDType | None = None) -> _CSRArray2D | csr_matrix: ... +@overload # mats: , format: "dia", dtype: +def block_diag(mats: _ToMatsDiagUnknown, format: _FmtDIA, dtype: _ToDType | None = None) -> _DIAArray | dia_matrix: ... +@overload # mats: , format: "dok", dtype: +def block_diag(mats: _ToMatsDiagUnknown, format: _FmtDOK, dtype: _ToDType | None = None) -> _DOKArray2D | dok_matrix: ... +@overload # mats: , format: "lil", dtype: +def block_diag(mats: _ToMatsDiagUnknown, format: _FmtLIL, dtype: _ToDType | None = None) -> _LILArray | lil_matrix: ... ### @overload # shape: T, format: , dtype: @@ -1827,7 +2026,7 @@ def random_array( *, density: float | npc.floating = 0.01, format: _FmtCOO = "coo", - dtype: npt.DTypeLike, + dtype: _ToDType, rng: onp.random.ToRNG | None = None, random_state: onp.random.ToRNG | None = None, data_sampler: _DataSampler | None = None, @@ -1838,7 +2037,7 @@ def random_array( *, density: float | npc.floating = 0.01, format: _FmtNonCOO, - dtype: npt.DTypeLike, + dtype: _ToDType, rng: onp.random.ToRNG | None = None, random_state: onp.random.ToRNG | None = None, data_sampler: _DataSampler | None = None, @@ -1982,7 +2181,7 @@ def random( n: opt.AnyInt, density: float | npc.floating = 0.01, format: _FmtCOO = "coo", - dtype: npt.DTypeLike | None = None, + dtype: _ToDType | None = None, rng: onp.random.ToRNG | None = None, data_rvs: _DataRVS | None = None, *, @@ -1995,7 +2194,7 @@ def random( density: float | npc.floating = 0.01, *, format: _FmtNonCOO, - dtype: npt.DTypeLike | None = None, + dtype: _ToDType | None = None, rng: onp.random.ToRNG | None = None, data_rvs: _DataRVS | None = None, random_state: onp.random.ToRNG | None = None, @@ -2006,7 +2205,7 @@ def random( n: opt.AnyInt, density: float | npc.floating, format: _FmtNonCOO, - dtype: npt.DTypeLike | None = None, + dtype: _ToDType | None = None, rng: onp.random.ToRNG | None = None, data_rvs: _DataRVS | None = None, *, @@ -2141,7 +2340,7 @@ def rand( n: opt.AnyInt, density: float | npc.floating = 0.01, format: _FmtCOO = "coo", - dtype: npt.DTypeLike | None = None, + dtype: _ToDType | None = None, rng: onp.random.ToRNG | None = None, *, random_state: onp.random.ToRNG | None = None, @@ -2153,7 +2352,7 @@ def rand( density: float | npc.floating = 0.01, *, format: _FmtNonCOO, - dtype: npt.DTypeLike | None = None, + dtype: _ToDType | None = None, rng: onp.random.ToRNG | None = None, random_state: onp.random.ToRNG | None = None, ) -> spmatrix: ... @@ -2163,7 +2362,7 @@ def rand( n: opt.AnyInt, density: float | npc.floating, format: _FmtNonCOO, - dtype: npt.DTypeLike | None = None, + dtype: _ToDType | None = None, rng: onp.random.ToRNG | None = None, *, random_state: onp.random.ToRNG | None = None, diff --git a/tests/sparse/test_construct.pyi b/tests/sparse/test_construct.pyi index ccee466a..eb0ea1d2 100644 --- a/tests/sparse/test_construct.pyi +++ b/tests/sparse/test_construct.pyi @@ -263,16 +263,43 @@ assert_type(sparse.hstack([csc_arr, csc_arr], format="lil", dtype=np.complex64), ### # block_array -assert_type(sparse.block_array([[any_mat]]), sparse.coo_array[ScalarType, tuple[int, int]]) -assert_type(sparse.block_array([[any_mat]], dtype=sctype), sparse.coo_array[ScalarType, tuple[int, int]]) -# bmat (legacy, `block_array` is preferred) -assert_type(sparse.bmat([[any_mat]]), sparse.coo_matrix[ScalarType]) -assert_type(sparse.bmat([[any_arr]]), sparse.coo_array[ScalarType, tuple[int, int]]) +assert_type(sparse.block_array([[bsr_mat]]), sparse.coo_matrix[ScalarType]) +assert_type(sparse.block_array([[bsr_arr]]), sparse.coo_array[ScalarType, tuple[int, int]]) +assert_type(sparse.block_array([[coo_arr]], dtype=int), sparse.coo_array[np.int64, tuple[int, int]]) +assert_type(sparse.block_array([[csr_arr]], dtype=sctype), sparse.csr_array[ScalarType, tuple[int, int]]) +assert_type(sparse.block_array([[lil_arr]], dtype=np.complex64), sparse.coo_array[np.complex64, tuple[int, int]]) +assert_type(sparse.block_array([[csr_arr]], format="lil"), sparse.lil_array[ScalarType]) +assert_type(sparse.block_array([[coo_arr, None]], format="bsr", dtype="float"), sparse.bsr_array[np.float64]) +assert_type(sparse.block_array([[dia_arr], [None]], format="dok", dtype=complex), sparse.dok_array[np.complex128]) +# NOTE: mypy reports false positive errors here, because of join vs. union approach +assert_type(sparse.block_array([[csr_arr, None], [None, dok_arr]], format="dia"), sparse.dia_array[ScalarType]) # type: ignore[assert-type] + +# bmat (legacy, `block_array` is p`referred) +assert_type(sparse.bmat([[bsr_mat]]), sparse.coo_matrix[ScalarType]) +assert_type(sparse.bmat([[dia_mat]], dtype=np.int_), sparse.coo_matrix[np.int_]) +assert_type(sparse.bmat([[bsr_arr]]), sparse.coo_array[ScalarType, tuple[int, int]]) +assert_type(sparse.bmat([[dia_arr]], dtype=np.int_), sparse.coo_array[np.int_, tuple[int, int]]) +# NOTE: mypy reports false positive errors here, because of join vs. union approach +assert_type(sparse.bmat([[csr_mat], [None]]), sparse.csr_matrix[ScalarType]) # type: ignore[assert-type,arg-type] +assert_type(sparse.bmat([[dok_mat], [None]], dtype=np.complex64), sparse.coo_matrix[np.complex64]) # type: ignore[assert-type,arg-type] +assert_type(sparse.bmat([[csr_arr], [None]]), sparse.csr_array[ScalarType, tuple[int, int]]) # type: ignore[assert-type,arg-type] +assert_type(sparse.bmat([[dok_arr], [None]], dtype=np.complex64), sparse.coo_array[np.complex64, tuple[int, int]]) # type: ignore[assert-type,arg-type] # block_diag -assert_type(sparse.block_diag([any_mat, any_mat]), sparse.coo_matrix[ScalarType]) assert_type(sparse.block_diag([any_arr, any_arr]), sparse.coo_array[ScalarType, tuple[int, int]]) -assert_type(sparse.block_diag([any_arr, any_mat]), sparse.coo_array[ScalarType, tuple[int, int]] | sparse.coo_matrix[ScalarType]) +assert_type(sparse.block_diag([any_arr, any_mat]), sparse.coo_matrix[ScalarType] | sparse.coo_array[ScalarType, tuple[int, int]]) +assert_type(sparse.block_diag([any_arr, any_arr], dtype="bool"), sparse.coo_array[np.bool_, tuple[int, int]]) +assert_type(sparse.block_diag([any_arr, any_arr], dtype=int), sparse.coo_array[np.int_, tuple[int, int]]) +assert_type( + sparse.block_diag([any_arr, any_mat], dtype=np.complex64), + sparse.coo_matrix[np.complex64] | sparse.coo_array[np.complex64, tuple[int, int]], +) +assert_type(sparse.block_diag([any_arr, any_arr], format="bsr"), sparse.bsr_array[ScalarType]) +assert_type(sparse.block_diag([any_arr, any_arr], format="csc", dtype=float), sparse.csc_array[np.float64]) +assert_type(sparse.block_diag([any_arr, any_arr], format="csr", dtype=complex), sparse.csr_array[np.complex128, tuple[int, int]]) +assert_type(sparse.block_diag([any_arr, any_arr], format="dia", dtype="int"), sparse.dia_array[np.int_]) +assert_type(sparse.block_diag([any_arr, any_arr], format="dok", dtype="bool"), sparse.dok_array[np.bool_, tuple[int, int]]) +assert_type(sparse.block_diag([any_arr, any_arr], format="lil", dtype=np.complex64), sparse.lil_array[np.complex64]) ### # random_array