Skip to content

Commit 581c5cb

Browse files
JulVandenBroeckjulvdb
JulVandenBroeck
authored and
julvdb
committed
Improved stubs for block_array and bmat 
1 parent eadadee commit 581c5cb

File tree

2 files changed

+177
-59
lines changed

2 files changed

+177
-59
lines changed

‎scipy-stubs/sparse/_construct.pyi‎

Lines changed: 159 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -76,7 +76,12 @@ _Offsets: TypeAlias = onp.ToInt | onp.ToInt1D
7676

7777
_DataRVS: TypeAlias = Callable[[int], onp.ArrayND[_Numeric]]
7878

79-
_ToBlocks: TypeAlias = Seq[Seq[_spbase[_SCT] | None]] | onp.ArrayND[np.object_]
79+
_ToBlocksArray: TypeAlias = onp.ArrayND[np.object_]
80+
_ToBlocksSPArray: TypeAlias = Seq[Seq[_SpArray2D[_SCT] | None]]
81+
_ToBlocksSPMatrix: TypeAlias = Seq[Seq[_SpMatrix[_SCT] | None]]
82+
_ToBlocksCanStack: TypeAlias = Seq[Seq[_CanStack[_T] | None]]
83+
_ToBlocksCanStackAs: TypeAlias = Seq[Seq[_CanStackAs[_SCT0, _T] | None]]
84+
_ToBlocksUnkown: TypeAlias = _ToBlocksSPArray | _ToBlocksArray
8085

8186
@type_check_only
8287
class _DataSampler(Protocol):
@@ -1378,19 +1383,19 @@ def kronsum(A: onp.ToComplex2D, B: onp.ToComplex2D, format: _Format | None = Non
13781383
@overload # sparray, format: <default>, dtype: <default>
13791384
def hstack(blocks: Seq[_CanStack[_T]], format: None = None, dtype: None = None) -> _T: ...
13801385
@overload # sparray, format: "bsr", dtype: <default>
1381-
def hstack(blocks: Seq[sparray[_SCT0]], format: _FmtBSR, dtype: None = None) -> _BSRArray[_SCT0]: ...
1386+
def hstack(blocks: Seq[_SpArray[_SCT]], format: _FmtBSR, dtype: None = None) -> _BSRArray[_SCT]: ...
13821387
@overload # sparray, format: "coo", dtype: <default>
1383-
def hstack(blocks: Seq[sparray[_SCT0]], format: _FmtCOO, dtype: None = None) -> _COOArray2D[_SCT0]: ...
1388+
def hstack(blocks: Seq[_SpArray[_SCT]], format: _FmtCOO, dtype: None = None) -> _COOArray2D[_SCT]: ...
13841389
@overload # sparray, format: "csc", dtype: <default>
1385-
def hstack(blocks: Seq[sparray[_SCT0]], format: _FmtCSC, dtype: None = None) -> _CSCArray[_SCT0]: ...
1390+
def hstack(blocks: Seq[_SpArray[_SCT]], format: _FmtCSC, dtype: None = None) -> _CSCArray[_SCT]: ...
13861391
@overload # sparray, format: "csr", dtype: <default>
1387-
def hstack(blocks: Seq[sparray[_SCT0]], format: _FmtCSR, dtype: None = None) -> _CSRArray2D[_SCT0]: ...
1392+
def hstack(blocks: Seq[_SpArray[_SCT]], format: _FmtCSR, dtype: None = None) -> _CSRArray2D[_SCT]: ...
13881393
@overload # sparray, format: "dia", dtype: <default>
1389-
def hstack(blocks: Seq[sparray[_SCT0]], format: _FmtDIA, dtype: None = None) -> _DIAArray[_SCT0]: ...
1394+
def hstack(blocks: Seq[_SpArray[_SCT]], format: _FmtDIA, dtype: None = None) -> _DIAArray[_SCT]: ...
13901395
@overload # sparray, format: "dok", dtype: <default>
1391-
def hstack(blocks: Seq[sparray[_SCT0]], format: _FmtDOK, dtype: None = None) -> _DOKArray2D[_SCT0]: ...
1396+
def hstack(blocks: Seq[_SpArray[_SCT]], format: _FmtDOK, dtype: None = None) -> _DOKArray2D[_SCT]: ...
13921397
@overload # sparray, format: "lil", dtype: <default>
1393-
def hstack(blocks: Seq[sparray[_SCT0]], format: _FmtLIL, dtype: None = None) -> _LILArray[_SCT0]: ...
1398+
def hstack(blocks: Seq[_SpArray[_SCT]], format: _FmtLIL, dtype: None = None) -> _LILArray[_SCT]: ...
13941399

13951400
#
13961401
@overload # sparray, format: <default>, dtype: bool-like
@@ -1509,19 +1514,19 @@ def hstack(blocks: Seq[_spbase], format: _Format, dtype: npt.DTypeLike | None =
15091514
@overload # sparray, format: <default>, dtype: <default>
15101515
def vstack(blocks: Seq[_CanStack[_T]], format: None = None, dtype: None = None) -> _T: ...
15111516
@overload # sparray, format: "bsr", dtype: <default>
1512-
def vstack(blocks: Seq[sparray[_SCT0]], format: _FmtBSR, dtype: None = None) -> _BSRArray[_SCT0]: ...
1517+
def vstack(blocks: Seq[_SpArray[_SCT]], format: _FmtBSR, dtype: None = None) -> _BSRArray[_SCT]: ...
15131518
@overload # sparray, format: "coo", dtype: <default>
1514-
def vstack(blocks: Seq[sparray[_SCT0]], format: _FmtCOO, dtype: None = None) -> _COOArray2D[_SCT0]: ...
1519+
def vstack(blocks: Seq[_SpArray[_SCT]], format: _FmtCOO, dtype: None = None) -> _COOArray2D[_SCT]: ...
15151520
@overload # sparray, format: "csc", dtype: <default>
1516-
def vstack(blocks: Seq[sparray[_SCT0]], format: _FmtCSC, dtype: None = None) -> _CSCArray[_SCT0]: ...
1521+
def vstack(blocks: Seq[_SpArray[_SCT]], format: _FmtCSC, dtype: None = None) -> _CSCArray[_SCT]: ...
15171522
@overload # sparray, format: "csr", dtype: <default>
1518-
def vstack(blocks: Seq[sparray[_SCT0]], format: _FmtCSR, dtype: None = None) -> _CSRArray2D[_SCT0]: ...
1523+
def vstack(blocks: Seq[_SpArray[_SCT]], format: _FmtCSR, dtype: None = None) -> _CSRArray2D[_SCT]: ...
15191524
@overload # sparray, format: "dia", dtype: <default>
1520-
def vstack(blocks: Seq[sparray[_SCT0]], format: _FmtDIA, dtype: None = None) -> _DIAArray[_SCT0]: ...
1525+
def vstack(blocks: Seq[_SpArray[_SCT]], format: _FmtDIA, dtype: None = None) -> _DIAArray[_SCT]: ...
15211526
@overload # sparray, format: "dok", dtype: <default>
1522-
def vstack(blocks: Seq[sparray[_SCT0]], format: _FmtDOK, dtype: None = None) -> _DOKArray2D[_SCT0]: ...
1527+
def vstack(blocks: Seq[_SpArray[_SCT]], format: _FmtDOK, dtype: None = None) -> _DOKArray2D[_SCT]: ...
15231528
@overload # sparray, format: "lil", dtype: <default>
1524-
def vstack(blocks: Seq[sparray[_SCT0]], format: _FmtLIL, dtype: None = None) -> _LILArray[_SCT0]: ...
1529+
def vstack(blocks: Seq[_SpArray[_SCT]], format: _FmtLIL, dtype: None = None) -> _LILArray[_SCT]: ...
15251530

15261531
#
15271532
@overload # sparray, format: <default>, dtype: bool-like
@@ -1636,52 +1641,152 @@ def vstack(blocks: Seq[sparray], format: _FmtLIL, dtype: npt.DTypeLike) -> _LILA
16361641
def vstack(blocks: Seq[_spbase], format: _Format, dtype: npt.DTypeLike | None = None) -> Incomplete: ...
16371642

16381643
###
1639-
# TODO(jorenham): Use `_CanStack` here, which requires a way to map matrix types to array types.
16401644
@overload # blocks: <known>, format: <default>, dtype: <default>
1641-
def block_array(blocks: _ToBlocks[_SCT], *, format: _FmtCOO | None = None, dtype: None = None) -> _COOArray2D[_SCT]: ...
1642-
@overload # blocks: <known>, format: <otherwise>, dtype: <default>
1643-
def block_array(blocks: _ToBlocks[_SCT], *, format: _FmtNonCOO, dtype: None = None) -> _SpArray2D[_SCT]: ...
1644-
1645-
#
1646-
@overload # blocks: <unknown>, format: <default>, dtype: <known>
1647-
def block_array(blocks: _ToBlocks, *, format: _FmtCOO | None = None, dtype: onp.ToDType[_SCT]) -> _COOArray2D[_SCT]: ...
1648-
@overload # blocks: <unknown>, format: <otherwise>, dtype: <known>
1649-
def block_array(blocks: _ToBlocks, *, format: _FmtNonCOO, dtype: onp.ToDType[_SCT]) -> _SpArray2D[_SCT]: ...
1650-
1651-
#
1652-
@overload # blocks: <unknown>, format: <default>, dtype: <unknown>
1653-
def block_array(blocks: _ToBlocks, *, format: _FmtCOO | None = None, dtype: npt.DTypeLike) -> _COOArray2D: ...
1654-
@overload # blocks: <unknown>, format: <otherwise>, dtype: <unknown>
1655-
def block_array(blocks: _ToBlocks, *, format: _FmtNonCOO, dtype: npt.DTypeLike) -> _SpArray2D: ...
1645+
def block_array(blocks: _ToBlocksCanStack[_T], format: None = None, dtype: None = None) -> _T: ...
1646+
@overload # blocks: <known>, format: "bsr", dtype: <default>
1647+
def block_array(blocks: _ToBlocksSPArray[_SCT], format: _FmtBSR, dtype: None = None) -> _BSRArray[_SCT]: ...
1648+
@overload # blocks: <known>, format: "coo", dtype: <default>
1649+
def block_array(blocks: _ToBlocksSPArray[_SCT], format: _FmtCOO, dtype: None = None) -> _COOArray2D[_SCT]: ...
1650+
@overload # blocks: <known>, format: "csc", dtype: <default>
1651+
def block_array(blocks: _ToBlocksSPArray[_SCT], format: _FmtCSC, dtype: None = None) -> _CSCArray[_SCT]: ...
1652+
@overload # blocks: <known>, format: "csr", dtype: <default>
1653+
def block_array(blocks: _ToBlocksSPArray[_SCT], format: _FmtCSR, dtype: None = None) -> _CSRArray2D[_SCT]: ...
1654+
@overload # blocks: <known>, format: "dia", dtype: <default>
1655+
def block_array(blocks: _ToBlocksSPArray[_SCT], format: _FmtDIA, dtype: None = None) -> _DIAArray[_SCT]: ...
1656+
@overload # blocks: <known>, format: "dok", dtype: <default>
1657+
def block_array(blocks: _ToBlocksSPArray[_SCT], format: _FmtDOK, dtype: None = None) -> _DOKArray2D[_SCT]: ...
1658+
@overload # blocks: <known>, format: "lil", dtype: <default>
1659+
def block_array(blocks: _ToBlocksSPArray[_SCT], format: _FmtLIL, dtype: None = None) -> _LILArray[_SCT]: ...
1660+
1661+
#
1662+
@overload # blocks: <known>, format: <default>, dtype: bool-like
1663+
def block_array(blocks: _ToBlocksCanStackAs[np.bool_, _T], format: None = None, *, dtype: onp.AnyBoolDType) -> _T: ...
1664+
@overload # blocks: <unknown>, format: "bsr", dtype: bool-like
1665+
def block_array(blocks: _ToBlocksUnkown, format: _FmtBSR, dtype: onp.AnyBoolDType) -> _BSRArray[np.bool_]: ...
1666+
@overload # blocks: <unknown>, format: "coo", dtype: bool-like
1667+
def block_array(blocks: _ToBlocksUnkown, format: _FmtCOO, dtype: onp.AnyBoolDType) -> _COOArray2D[np.bool_]: ...
1668+
@overload # blocks: <unknown>, format: "csc", dtype: bool-like
1669+
def block_array(blocks: _ToBlocksUnkown, format: _FmtCSC, dtype: onp.AnyBoolDType) -> _CSCArray[np.bool_]: ...
1670+
@overload # blocks: <unknown>, format: "csr", dtype: bool-like
1671+
def block_array(blocks: _ToBlocksUnkown, format: _FmtCSR, dtype: onp.AnyBoolDType) -> _CSRArray2D[np.bool_]: ...
1672+
@overload # blocks: <unknown>, format: "dia", dtype: bool-like
1673+
def block_array(blocks: _ToBlocksUnkown, format: _FmtDIA, dtype: onp.AnyBoolDType) -> _DIAArray[np.bool_]: ...
1674+
@overload # blocks: <unknown>, format: "dok", dtype: bool-like
1675+
def block_array(blocks: _ToBlocksUnkown, format: _FmtDOK, dtype: onp.AnyBoolDType) -> _DOKArray2D[np.bool_]: ...
1676+
@overload # blocks: <unknown>, format: "lil", dtype: bool-like
1677+
def block_array(blocks: _ToBlocksUnkown, format: _FmtLIL, dtype: onp.AnyBoolDType) -> _LILArray[np.bool_]: ...
1678+
1679+
#
1680+
@overload # blocks: <known>, format: <default>, dtype: int-like
1681+
def block_array(blocks: _ToBlocksCanStackAs[np.int_, _T], format: None = None, *, dtype: onp.AnyIntDType) -> _T: ...
1682+
@overload # blocks: <unknown>, format: "bsr", dtype: int-like
1683+
def block_array(blocks: _ToBlocksUnkown, format: _FmtBSR, dtype: onp.AnyIntDType) -> _BSRArray[np.int_]: ...
1684+
@overload # blocks: <unknown>, format: "coo", dtype: int-like
1685+
def block_array(blocks: _ToBlocksUnkown, format: _FmtCOO, dtype: onp.AnyIntDType) -> _COOArray2D[np.int_]: ...
1686+
@overload # blocks: <unknown>, format: "csc", dtype: int-like
1687+
def block_array(blocks: _ToBlocksUnkown, format: _FmtCSC, dtype: onp.AnyIntDType) -> _CSCArray[np.int_]: ...
1688+
@overload # blocks: <unknown>, format: "csr", dtype: int-like
1689+
def block_array(blocks: _ToBlocksUnkown, format: _FmtCSR, dtype: onp.AnyIntDType) -> _CSRArray2D[np.int_]: ...
1690+
@overload # blocks: <unknown>, format: "dia", dtype: int-like
1691+
def block_array(blocks: _ToBlocksUnkown, format: _FmtDIA, dtype: onp.AnyIntDType) -> _DIAArray[np.int_]: ...
1692+
@overload # blocks: <unknown>, format: "dok", dtype: int-like
1693+
def block_array(blocks: _ToBlocksUnkown, format: _FmtDOK, dtype: onp.AnyIntDType) -> _DOKArray2D[np.int_]: ...
1694+
@overload # blocks: <unknown>, format: "lil", dtype: int-like
1695+
def block_array(blocks: _ToBlocksUnkown, format: _FmtLIL, dtype: onp.AnyIntDType) -> _LILArray[np.int_]: ...
1696+
1697+
#
1698+
@overload # blocks: <known>, format: <default>, dtype: float64-like
1699+
def block_array(blocks: _ToBlocksCanStackAs[np.float64, _T], format: None = None, *, dtype: onp.AnyFloat64DType) -> _T: ...
1700+
@overload # blocks: <unknown>, format: "bsr", dtype: float64-like
1701+
def block_array(blocks: _ToBlocksUnkown, format: _FmtBSR, dtype: onp.AnyFloat64DType) -> _BSRArray[np.float64]: ...
1702+
@overload # blocks: <unknown>, format: "coo", dtype: float64-like
1703+
def block_array(blocks: _ToBlocksUnkown, format: _FmtCOO, dtype: onp.AnyFloat64DType) -> _COOArray2D[np.float64]: ...
1704+
@overload # blocks: <unknown>, format: "csc", dtype: float64-like
1705+
def block_array(blocks: _ToBlocksUnkown, format: _FmtCSC, dtype: onp.AnyFloat64DType) -> _CSCArray[np.float64]: ...
1706+
@overload # blocks: <unknown>, format: "csr", dtype: float64-like
1707+
def block_array(blocks: _ToBlocksUnkown, format: _FmtCSR, dtype: onp.AnyFloat64DType) -> _CSRArray2D[np.float64]: ...
1708+
@overload # blocks: <unknown>, format: "dia", dtype: float64-like
1709+
def block_array(blocks: _ToBlocksUnkown, format: _FmtDIA, dtype: onp.AnyFloat64DType) -> _DIAArray[np.float64]: ...
1710+
@overload # blocks: <unknown>, format: "dok", dtype: float64-like
1711+
def block_array(blocks: _ToBlocksUnkown, format: _FmtDOK, dtype: onp.AnyFloat64DType) -> _DOKArray2D[np.float64]: ...
1712+
@overload # blocks: <unknown>, format: "lil", dtype: float64-like
1713+
def block_array(blocks: _ToBlocksUnkown, format: _FmtLIL, dtype: onp.AnyFloat64DType) -> _LILArray[np.float64]: ...
1714+
1715+
#
1716+
@overload # blocks: <known>, format: <default>, dtype: complex128-like
1717+
def block_array(blocks: _ToBlocksCanStackAs[np.complex128, _T], format: None = None, *, dtype: onp.AnyComplex128DType) -> _T: ...
1718+
@overload # blocks: <unknown>, format: "bsr", dtype: complex128-like
1719+
def block_array(blocks: _ToBlocksUnkown, format: _FmtBSR, dtype: onp.AnyComplex128DType) -> _BSRArray[np.complex128]: ...
1720+
@overload # blocks: <unknown>, format: "coo", dtype: complex128-like
1721+
def block_array(blocks: _ToBlocksUnkown, format: _FmtCOO, dtype: onp.AnyComplex128DType) -> _COOArray2D[np.complex128]: ...
1722+
@overload # blocks: <unknown>, format: "csc", dtype: complex128-like
1723+
def block_array(blocks: _ToBlocksUnkown, format: _FmtCSC, dtype: onp.AnyComplex128DType) -> _CSCArray[np.complex128]: ...
1724+
@overload # blocks: <unknown>, format: "csr", dtype: complex128-like
1725+
def block_array(blocks: _ToBlocksUnkown, format: _FmtCSR, dtype: onp.AnyComplex128DType) -> _CSRArray2D[np.complex128]: ...
1726+
@overload # blocks: <unknown>, format: "dia", dtype: complex128-like
1727+
def block_array(blocks: _ToBlocksUnkown, format: _FmtDIA, dtype: onp.AnyComplex128DType) -> _DIAArray[np.complex128]: ...
1728+
@overload # blocks: <unknown>, format: "dok", dtype: complex128-like
1729+
def block_array(blocks: _ToBlocksUnkown, format: _FmtDOK, dtype: onp.AnyComplex128DType) -> _DOKArray2D[np.complex128]: ...
1730+
@overload # blocks: <unknown>, format: "lil", dtype: complex128-like
1731+
def block_array(blocks: _ToBlocksUnkown, format: _FmtLIL, dtype: onp.AnyComplex128DType) -> _LILArray[np.complex128]: ...
1732+
1733+
#
1734+
@overload # blocks: <known>, format: <default>, dtype: <known>
1735+
def block_array(blocks: _ToBlocksCanStackAs[_SCT0, _T], format: None = None, *, dtype: onp.ToDType[_SCT0]) -> _T: ...
1736+
@overload # blocks: <unknown>, format: "bsr", dtype: <known>
1737+
def block_array(blocks: _ToBlocksUnkown, format: _FmtBSR, dtype: onp.ToDType[_SCT0]) -> _BSRArray[_SCT0]: ...
1738+
@overload # blocks: <unknown>, format: "coo", dtype: <known>
1739+
def block_array(blocks: _ToBlocksUnkown, format: _FmtCOO, dtype: onp.ToDType[_SCT0]) -> _COOArray2D[_SCT0]: ...
1740+
@overload # blocks: <unknown>, format: "csc", dtype: <known>
1741+
def block_array(blocks: _ToBlocksUnkown, format: _FmtCSC, dtype: onp.ToDType[_SCT0]) -> _CSCArray[_SCT0]: ...
1742+
@overload # blocks: <unknown>, format: "csr", dtype: <known>
1743+
def block_array(blocks: _ToBlocksUnkown, format: _FmtCSR, dtype: onp.ToDType[_SCT0]) -> _CSRArray2D[_SCT0]: ...
1744+
@overload # blocks: <unknown>, format: "dia", dtype: <known>
1745+
def block_array(blocks: _ToBlocksUnkown, format: _FmtDIA, dtype: onp.ToDType[_SCT0]) -> _DIAArray[_SCT0]: ...
1746+
@overload # blocks: <unknown>, format: "dok", dtype: <known>
1747+
def block_array(blocks: _ToBlocksUnkown, format: _FmtDOK, dtype: onp.ToDType[_SCT0]) -> _DOKArray2D[_SCT0]: ...
1748+
@overload # blocks: <unknown>, format: "lil", dtype: <known>
1749+
def block_array(blocks: _ToBlocksUnkown, format: _FmtLIL, dtype: onp.ToDType[_SCT0]) -> _LILArray[_SCT0]: ...
1750+
1751+
#
1752+
@overload # blocks: <known>, format: <default>, dtype: <unknown>
1753+
def block_array(blocks: _ToBlocksCanStackAs[Any, _T], format: None = None, *, dtype: npt.DTypeLike) -> _T: ...
1754+
@overload # blocks: <unknown>, format: "bsr", dtype: <unknown>
1755+
def block_array(blocks: _ToBlocksUnkown, format: _FmtBSR, dtype: npt.DTypeLike) -> _BSRArray: ...
1756+
@overload # blocks: <unknown>, format: "coo", dtype: <unknown>
1757+
def block_array(blocks: _ToBlocksUnkown, format: _FmtCOO, dtype: npt.DTypeLike) -> _COOArray2D: ...
1758+
@overload # blocks: <unknown>, format: "csc", dtype: <unknown>
1759+
def block_array(blocks: _ToBlocksUnkown, format: _FmtCSC, dtype: npt.DTypeLike) -> _CSCArray: ...
1760+
@overload # blocks: <unknown>, format: "csr", dtype: <unknown>
1761+
def block_array(blocks: _ToBlocksUnkown, format: _FmtCSR, dtype: npt.DTypeLike) -> _CSRArray2D: ...
1762+
@overload # blocks: <unknown>, format: "dia", dtype: <unknown>
1763+
def block_array(blocks: _ToBlocksUnkown, format: _FmtDIA, dtype: npt.DTypeLike) -> _DIAArray: ...
1764+
@overload # blocks: <unknown>, format: "dok", dtype: <unknown>
1765+
def block_array(blocks: _ToBlocksUnkown, format: _FmtDOK, dtype: npt.DTypeLike) -> _DOKArray2D: ...
1766+
@overload # blocks: <unknown>, format: "lil", dtype: <unknown>
1767+
def block_array(blocks: _ToBlocksUnkown, format: _FmtLIL, dtype: npt.DTypeLike) -> _LILArray: ...
16561768

16571769
###
1658-
# TODO(jorenham): Use `_CanStack` here, which requires a way to map array types to matrix types.
1659-
@overload # blocks: <array, known>, format: <default>, dtype: <default>
1660-
def bmat(blocks: Seq[Seq[sparray[_SCT]]], format: _FmtCOO | None = None, dtype: None = None) -> _COOArray2D[_SCT]: ...
1661-
@overload # blocks: <array, known>, format: <otherwise>, dtype: <default>
1662-
def bmat(blocks: Seq[Seq[sparray[_SCT]]], format: _FmtNonCOO, dtype: None = None) -> _SpArray2D[_SCT]: ...
1663-
1664-
#
1665-
@overload # blocks: <matrix, known>, format: <default>, dtype: <default>
1666-
def bmat(blocks: Seq[Seq[spmatrix[_SCT]]], format: _FmtCOO | None = None, dtype: None = None) -> coo_matrix[_SCT]: ...
1667-
@overload # blocks: <matrix, known>, format: <otherwise>, dtype: <default>
1668-
def bmat(blocks: Seq[Seq[spmatrix[_SCT]]], format: _FmtNonCOO, dtype: None = None) -> spmatrix[_SCT]: ...
1770+
@overload # blocks: <known>, format: <default>, dtype: <default>
1771+
def bmat(blocks: _ToBlocksCanStack[_T], format: None = None, dtype: None = None) -> _T: ...
1772+
@overload # blocks: <matrix>, format: <otherwise>, dtype: <default>
1773+
def bmat(blocks:_ToBlocksSPMatrix[_SCT], format: _Format, dtype: None = None) -> _SpMatrix[_SCT]: ...
16691774

16701775
#
1671-
@overload # blocks: <unknown, unknown>, format: <default>, dtype: <known>
1672-
def bmat(
1673-
blocks: _ToBlocks, format: _FmtCOO | None = None, *, dtype: onp.ToDType[_SCT]
1674-
) -> _COOArray2D[_SCT] | coo_matrix[_SCT]: ...
1675-
@overload # blocks: <unknown, unknown>, format: <otherwise>, dtype: <known>
1676-
def bmat(blocks: _ToBlocks, format: _FmtNonCOO, dtype: onp.ToDType[_SCT]) -> _SpBase2D[_SCT]: ...
1776+
@overload # blocks: <known>, format: <default>, dtype: <known>
1777+
def bmat(blocks: _ToBlocksCanStackAs[_SCT0, _T], format: None = None, *, dtype: onp.ToDType[_SCT0]) -> _T: ...
1778+
@overload # blocks: <matrix>, format: <otherwise>, dtype: <known>
1779+
def bmat(blocks: _ToBlocksSPMatrix[_SCT0], format: _Format, dtype: onp.ToDType[_SCT0]) -> _SpMatrix[_SCT0]: ...
1780+
@overload # blocks: <unknown>, format: <otherwise>, dtype: <known>
1781+
def bmat(blocks: _ToBlocksUnkown, format: _Format, dtype: onp.ToDType[_SCT0]) -> _SpMatrix[_SCT0] | _SpArray2D[_SCT0]: ...
16771782

16781783
#
1679-
@overload # blocks: <unknown, unknown>, format: <default>, dtype: <unknown>
1680-
def bmat(
1681-
blocks: _ToBlocks[_SCT], format: _FmtCOO | None = None, dtype: npt.DTypeLike | None = None
1682-
) -> _COOArray2D[_SCT] | coo_matrix[_SCT]: ...
1683-
@overload # blocks: <unknown, unknown>, format: <otherwise>, dtype: <unknown>
1684-
def bmat(blocks: _ToBlocks[_SCT], format: _FmtNonCOO, dtype: npt.DTypeLike) -> _SpBase2D[_SCT]: ...
1784+
@overload # blocks: <known>, format: <default>, dtype: <unknown>
1785+
def bmat(blocks: _ToBlocksCanStackAs[Any, _T], format: None = None, *, dtype: npt.DTypeLike) -> _T: ...
1786+
@overload # blocks: <matrix>, format: <otherwise>, dtype: <unknown>
1787+
def bmat(blocks: _ToBlocksSPMatrix, format: _Format, *, dtype: npt.DTypeLike) -> _SpMatrix: ...
1788+
@overload # blocks: <unknown>, format: <otherwise>, dtype: <unknown>
1789+
def bmat(blocks: _ToBlocksUnkown, format: _Format, *, dtype: npt.DTypeLike) -> _SpMatrix | _SpArray2D: ...
16851790

16861791
###
16871792
# TODO(jorenham): Add support for non-COO formats.

0 commit comments

Comments
 (0)