diff --git a/scipy-stubs/_typing.pyi b/scipy-stubs/_typing.pyi index 3c933602..f980d182 100644 --- a/scipy-stubs/_typing.pyi +++ b/scipy-stubs/_typing.pyi @@ -25,6 +25,7 @@ __all__ = [ "FileModeRWA", "FileName", "NanPolicy", + "OrderCF", "OrderKACF", "Seed", "Untyped", @@ -72,6 +73,7 @@ AnyShape: TypeAlias = op.CanIndex | Sequence[op.CanIndex] RNG: TypeAlias = np.random.BitGenerator | np.random.Generator | np.random.RandomState Seed: TypeAlias = RNG | np.random.SeedSequence | onp.ToJustInt | onp.ToJustIntND ByteOrder: TypeAlias = Literal["S", "<", "little", ">", "big", "=", "native", "|", "I"] +OrderCF: TypeAlias = Literal["C", "F"] OrderKACF: TypeAlias = Literal["K", "A", "C", "F"] Casting: TypeAlias = Literal["no", "equiv", "safe", "same_kind", "unsafe"] CorrelateMode: TypeAlias = Literal["valid", "same", "full"] diff --git a/scipy-stubs/sparse/_base.pyi b/scipy-stubs/sparse/_base.pyi index f9fcee04..a0119b80 100644 --- a/scipy-stubs/sparse/_base.pyi +++ b/scipy-stubs/sparse/_base.pyi @@ -1,38 +1,39 @@ +# mypy: disable-error-code="misc" # pyright: reportUnannotatedClassAttribute=false import abc from collections.abc import Iterator, Sequence -from typing import Any, Final, Generic, Literal, TypeAlias, overload -from typing_extensions import Self, TypeAliasType, TypeIs, TypeVar +from typing import Final, Generic, Literal, TypeAlias, overload +from typing_extensions import Self, TypeIs, TypeVar import numpy as np import optype as op import optype.numpy as onp import optype.typing as opt -from scipy._typing import Untyped +from scipy._typing import Casting, OrderCF, Untyped +from ._bsr import bsr_array, bsr_matrix +from ._coo import coo_array, coo_matrix +from ._csc import csc_array, csc_matrix +from ._csr import csr_array, csr_matrix +from ._dia import dia_array, dia_matrix +from ._dok import dok_array, dok_matrix +from ._lil import lil_array, lil_matrix from ._matrix import spmatrix as spmatrix +from ._typing import Float, Int, Scalar, SPFormat, ToShape __all__ = ["SparseEfficiencyWarning", "SparseWarning", "issparse", "isspmatrix", "sparray"] -# NOTE: This is roughly speaking equivalent to `[u]int8 | [u]int16 | [u]int32 | [u]int64` (on most modern platforms) -_Int: TypeAlias = np.integer[Any] -# NOTE: NOT equivalent to `floating`! It's considered invalid to use `float16` in `scipy.sparse`. -_Float: TypeAlias = np.float32 | np.float64 | np.longdouble -# NOTE: This is (almost always) equivalent to `complex64 | complex128 | clongdouble`. -_Complex: TypeAlias = np.complexfloating[Any, Any] -# NOTE: Roughly speaking, this is equivalent to `number[Any] - float16` (the `-` denotes the set difference analogue for types) -# NOTE: This (almost always) matches `._sputils.supported_dtypes` -# NOTE: The `TypeAliasType` is used to avoid long error messages. -_Scalar = TypeAliasType("_Scalar", np.bool_ | _Int | _Float | _Complex) - _T = TypeVar("_T") -_SCT_co = TypeVar("_SCT_co", bound=_Scalar, default=_Scalar, covariant=True) +_OutT = TypeVar("_OutT", bound=onp.ArrayND) +_SCT = TypeVar("_SCT", bound=np.generic) +_SCT_co = TypeVar("_SCT_co", bound=Scalar, default=Scalar, covariant=True) + +_Matrix: TypeAlias = np.matrix[tuple[int, int], np.dtype[_SCT]] -_ToShape: TypeAlias = tuple[op.CanIndex] | tuple[op.CanIndex, op.CanIndex] _ToSparseFromPy: TypeAlias = Sequence[Sequence[_T]] | Sequence[_T] _ToSparseFromArrayLike: TypeAlias = onp.CanArrayND[_SCT_co] | _ToSparseFromPy[_SCT_co] -_Format: TypeAlias = Literal["bsr", "coo", "csc", "csr", "dia", "dok", "lil"] +_Real: TypeAlias = np.bool_ | Int | Float ### @@ -42,8 +43,7 @@ class SparseWarning(Warning): ... class SparseFormatWarning(SparseWarning): ... class SparseEfficiencyWarning(SparseWarning): ... -# TODO(jorenham): Generic `shape` property -# TODO(jorenham): Generic `data` attribute +# TODO(jorenham): Generic `shape` and `data` class _spbase(Generic[_SCT_co]): __array_priority__: float = 10.1 maxprint: Final[int | None] @@ -61,7 +61,7 @@ class _spbase(Generic[_SCT_co]): # NOTE: At runtime this isn't abstract, but returns `und` instead. @property @abc.abstractmethod - def format(self, /) -> _Format: ... + def format(self, /) -> SPFormat: ... # @property @@ -73,7 +73,7 @@ class _spbase(Generic[_SCT_co]): # NOTE: In `scipy>=1.15.0` the `maxprint` param will become keyword-only. @overload # shape - def __init__(self: _spbase[np.float64], /, arg1: _ToShape, maxprint: int | None = 50) -> None: ... + def __init__(self: _spbase[np.float64], /, arg1: ToShape, maxprint: int | None = 50) -> None: ... @overload # sparse def __init__(self, /, arg1: _spbase[_SCT_co], maxprint: int | None = 50) -> None: ... @overload # dense array-like @@ -108,13 +108,38 @@ class _spbase(Generic[_SCT_co]): # def __bool__(self, /) -> bool: ... - def __iter__(self, /) -> Iterator[Untyped]: ... # - def __lt__(self, other: Untyped, /) -> Untyped: ... - def __gt__(self, other: Untyped, /) -> Untyped: ... - def __le__(self, other: Untyped, /) -> Untyped: ... - def __ge__(self, other: Untyped, /) -> Untyped: ... + @overload # {csc,csr}_array + def __iter__(self: csc_array | csr_array, /) -> Iterator[csr_array[_SCT_co]]: ... + @overload # {csc,csr}_matrix + def __iter__(self: csc_matrix | csr_matrix, /) -> Iterator[csr_matrix[_SCT_co]]: ... + @overload # dok_array + def __iter__(self: dok_array, /) -> Iterator[dok_array[_SCT_co]]: ... + @overload # dok_matrix + def __iter__(self: dok_matrix, /) -> Iterator[dok_matrix[_SCT_co]]: ... + @overload # dok_array + def __iter__(self: lil_array, /) -> Iterator[lil_array[_SCT_co]]: ... + @overload # dok_array + def __iter__(self: lil_matrix, /) -> Iterator[lil_matrix[_SCT_co]]: ... + + # + @overload + def __lt__(self: sparray, other: _spbase[_Real] | onp.ToFloat, /) -> csr_array[np.bool_]: ... + @overload + def __lt__(self: spmatrix, other: _spbase[_Real], /) -> csr_matrix[np.bool_]: ... + @overload + def __gt__(self: sparray, other: _spbase[_Real] | onp.ToFloat, /) -> csr_array[np.bool_]: ... + @overload + def __gt__(self: spmatrix, other: _spbase[_Real], /) -> csr_matrix[np.bool_]: ... + @overload + def __le__(self: sparray, other: _spbase[_Real] | onp.ToFloat, /) -> csr_array[np.bool_]: ... + @overload + def __le__(self: spmatrix, other: _spbase[_Real], /) -> csr_matrix[np.bool_]: ... + @overload + def __ge__(self: sparray, other: _spbase[_Real] | onp.ToFloat, /) -> csr_array[np.bool_]: ... + @overload + def __ge__(self: spmatrix, other: _spbase[_Real], /) -> csr_matrix[np.bool_]: ... # def __neg__(self, /) -> Self: ... @@ -131,14 +156,14 @@ class _spbase(Generic[_SCT_co]): def __matmul__(self, other: Untyped, /) -> Untyped: ... def __rmatmul__(self, other: Untyped, /) -> Untyped: ... def __truediv__(self, other: Untyped, /) -> Untyped: ... - def __div__(self, other: Untyped, /) -> Untyped: ... + def __div__(self, other: Untyped, /) -> Untyped: ... # que? def __pow__(self, other: Untyped, /) -> Untyped: ... # - def nonzero(self, /) -> Untyped: ... + def nonzero(self, /) -> tuple[onp.Array1D[np.int32], onp.Array1D[np.int32]]: ... def count_nonzero(self, /) -> int: ... def conjugate(self, /, copy: bool = True) -> Self: ... - conj = conjugate + def conj(self, /, copy: bool = True) -> Self: ... def transpose(self, /, axes: Untyped | None = None, copy: bool = False) -> Self: ... # @@ -149,31 +174,80 @@ class _spbase(Generic[_SCT_co]): def power(self, /, n: Untyped, dtype: Untyped | None = None) -> Self: ... # - def diagonal(self, /, k: int = 0) -> Untyped: ... - def trace(self, /, offset: int = 0) -> Untyped: ... - def sum(self, /, axis: Untyped | None = None, dtype: Untyped | None = None, out: Untyped | None = None) -> Untyped: ... - def mean(self, /, axis: Untyped | None = None, dtype: Untyped | None = None, out: Untyped | None = None) -> Untyped: ... + def diagonal(self, /, k: int = 0) -> onp.Array1D[_SCT_co]: ... # only if 2-d + def trace(self, /, offset: int = 0) -> _SCT_co: ... + + # + def sum(self, /, axis: op.CanIndex | None = None, dtype: Untyped | None = None, out: Untyped | None = None) -> Untyped: ... + def mean(self, /, axis: op.CanIndex | None = None, dtype: Untyped | None = None, out: Untyped | None = None) -> Untyped: ... # def copy(self, /) -> Self: ... def reshape(self, /, *args: Untyped, **kwargs: Untyped) -> Self: ... - def astype(self, /, dtype: Untyped, casting: str = "unsafe", copy: bool = True) -> Untyped: ... + def astype(self, /, dtype: Untyped, casting: Casting = "unsafe", copy: bool = True) -> Untyped: ... def asformat(self, /, format: Untyped, copy: bool = False) -> Untyped: ... # - def todense(self, /, order: Untyped | None = None, out: Untyped | None = None) -> Untyped: ... - def toarray(self, /, order: Untyped | None = None, out: Untyped | None = None) -> Untyped: ... - def tocsr(self, /, copy: bool = False) -> Untyped: ... - def todok(self, /, copy: bool = False) -> Untyped: ... - def tocoo(self, /, copy: bool = False) -> Untyped: ... - def tolil(self, /, copy: bool = False) -> Untyped: ... - def todia(self, /, copy: bool = False) -> Untyped: ... - def tobsr(self, /, blocksize: tuple[int, int] | None = None, copy: bool = False) -> Untyped: ... - def tocsc(self, /, copy: bool = False) -> Untyped: ... + @overload # self: spmatrix, out: None + def todense(self: spmatrix, /, order: OrderCF | None = None, out: None = None) -> _Matrix[_SCT_co]: ... + @overload # self: spmatrix, out: array (positional) + def todense(self: spmatrix, /, order: OrderCF | None, out: onp.ArrayND[_SCT]) -> _Matrix[_SCT]: ... + @overload # self: spmatrix, out: array (keyword) + def todense(self: spmatrix, /, order: OrderCF | None = None, *, out: onp.ArrayND[_SCT]) -> _Matrix[_SCT]: ... + @overload # out: None + def todense(self, /, order: OrderCF | None = None, out: None = None) -> onp.Array2D[_SCT_co]: ... + @overload # out: array (positional) + def todense(self, /, order: OrderCF | None, out: _OutT) -> _OutT: ... + @overload # out: array (keyword) + def todense(self, /, order: OrderCF | None = None, *, out: _OutT) -> _OutT: ... + + # + @overload # out: None + def toarray(self, /, order: OrderCF | None = None, out: None = None) -> onp.Array2D[_SCT_co]: ... + @overload # out: array (positional) + def toarray(self, /, order: OrderCF | None, out: _OutT) -> _OutT: ... + @overload # out: array (keyword) + def toarray(self, /, order: OrderCF | None = None, *, out: _OutT) -> _OutT: ... + + # + @overload + def tobsr(self: sparray, /, blocksize: tuple[int, int] | None = None, copy: bool = False) -> bsr_array[_SCT_co]: ... + @overload + def tobsr(self: spmatrix, /, blocksize: tuple[int, int] | None = None, copy: bool = False) -> bsr_matrix[_SCT_co]: ... + # + @overload + def tocoo(self: sparray, /, copy: bool = False) -> coo_array[_SCT_co]: ... + @overload + def tocoo(self: spmatrix, /, copy: bool = False) -> coo_matrix[_SCT_co]: ... + # + @overload + def tocsc(self: sparray, /, copy: bool = False) -> csc_array[_SCT_co]: ... + @overload + def tocsc(self: spmatrix, /, copy: bool = False) -> csc_matrix[_SCT_co]: ... + # + @overload + def tocsr(self: sparray, /, copy: bool = False) -> csr_array[_SCT_co]: ... + @overload + def tocsr(self: spmatrix, /, copy: bool = False) -> csr_matrix[_SCT_co]: ... + # + @overload + def todia(self: sparray, /, copy: bool = False) -> dia_array[_SCT_co]: ... + @overload + def todia(self: spmatrix, /, copy: bool = False) -> dia_matrix[_SCT_co]: ... + # + @overload + def todok(self: sparray, /, copy: bool = False) -> dok_array[_SCT_co]: ... + @overload + def todok(self: spmatrix, /, copy: bool = False) -> dok_matrix[_SCT_co]: ... + # + @overload + def tolil(self: sparray, /, copy: bool = False) -> lil_array[_SCT_co]: ... + @overload + def tolil(self: spmatrix, /, copy: bool = False) -> lil_matrix[_SCT_co]: ... # def resize(self, /, shape: tuple[op.CanIndex, op.CanIndex]) -> None: ... - def setdiag(self, /, values: Untyped, k: int = 0) -> None: ... + def setdiag(self, /, values: onp.ToComplex1D, k: int = 0) -> None: ... class sparray: ... diff --git a/scipy-stubs/sparse/_bsr.pyi b/scipy-stubs/sparse/_bsr.pyi index f6c24a22..17892aed 100644 --- a/scipy-stubs/sparse/_bsr.pyi +++ b/scipy-stubs/sparse/_bsr.pyi @@ -1,41 +1,130 @@ -from typing import Literal -from typing_extensions import override +from collections.abc import Sequence +from typing import Any, Generic, Literal, TypeAlias, overload +from typing_extensions import TypeIs, TypeVar, override -import optype as op -from scipy._typing import Untyped -from ._base import sparray +import numpy as np +import optype.numpy as onp +import optype.typing as opt +from ._base import _spbase, sparray from ._compressed import _cs_matrix from ._data import _minmax_mixin from ._matrix import spmatrix +from ._typing import Int, Scalar, ToShape2D __all__ = ["bsr_array", "bsr_matrix", "isspmatrix_bsr"] -# TODO(jorenham): generic dtype -class _bsr_base(_cs_matrix, _minmax_mixin): - data: Untyped - indices: Untyped - indptr: Untyped +_T = TypeVar("_T") +_SCT = TypeVar("_SCT", bound=Scalar, default=Any) + +_ToDType: TypeAlias = type[_SCT] | np.dtype[_SCT] | onp.HasDType[np.dtype[_SCT]] +_ToMatrix: TypeAlias = _spbase[_SCT] | onp.CanArrayND[_SCT] | Sequence[onp.CanArrayND[_SCT]] | _ToMatrixPy[_SCT] +_ToMatrixPy: TypeAlias = Sequence[_T] | Sequence[Sequence[_T]] + +_ToData2: TypeAlias = tuple[onp.ArrayND[_SCT], onp.ArrayND[Int]] +_ToData3: TypeAlias = tuple[onp.ArrayND[_SCT], onp.ArrayND[Int], onp.ArrayND[Int]] +_ToData: TypeAlias = _ToData2[_SCT] | _ToData3[_SCT] + +### + +class _bsr_base(_cs_matrix[_SCT], _minmax_mixin[_SCT], Generic[_SCT]): + data: onp.Array3D[_SCT] @property @override def format(self, /) -> Literal["bsr"]: ... - - # + @property + @override + def ndim(self, /) -> Literal[2]: ... + @property + @override + def shape(self, /) -> tuple[int, int]: ... @property def blocksize(self, /) -> tuple[int, int]: ... # + @overload # matrix-like (known dtype), dtype: None + def __init__( + self, + /, + arg1: _ToMatrix[_SCT] | _ToData[_SCT], + shape: ToShape2D | None = None, + dtype: None = None, + copy: bool = False, + blocksize: tuple[int, int] | None = None, + ) -> None: ... + @overload # 2-d shape-like, dtype: None + def __init__( + self: _bsr_base[np.float64], + /, + arg1: ToShape2D, + shape: None = None, + dtype: None = None, + copy: bool = False, + blocksize: tuple[int, int] | None = None, + ) -> None: ... + @overload # matrix-like builtins.bool, dtype: type[bool] | None + def __init__( + self: _bsr_base[np.bool_], + /, + arg1: _ToMatrixPy[bool], + shape: ToShape2D | None = None, + dtype: onp.AnyBoolDType | None = None, + copy: bool = False, + blocksize: tuple[int, int] | None = None, + ) -> None: ... + @overload # matrix-like builtins.int, dtype: type[int] | None + def __init__( + self: _bsr_base[np.int_], + /, + arg1: _ToMatrixPy[opt.JustInt], + shape: ToShape2D | None = None, + dtype: type[opt.JustInt] | onp.AnyIntPDType | None = None, + copy: bool = False, + blocksize: tuple[int, int] | None = None, + ) -> None: ... + @overload # matrix-like builtins.float, dtype: type[float] | None + def __init__( + self: _bsr_base[np.float64], + /, + arg1: _ToMatrixPy[opt.Just[float]], + shape: ToShape2D | None = None, + dtype: type[opt.Just[float]] | onp.AnyFloat64DType | None = None, + copy: bool = False, + blocksize: tuple[int, int] | None = None, + ) -> None: ... + @overload # matrix-like builtins.complex, dtype: type[complex] | None + def __init__( + self: _bsr_base[np.complex128], + /, + arg1: _ToMatrixPy[opt.Just[complex]], + shape: ToShape2D | None = None, + dtype: type[opt.Just[complex]] | onp.AnyComplex128DType | None = None, + copy: bool = False, + blocksize: tuple[int, int] | None = None, + ) -> None: ... + @overload # dtype: (positional) + def __init__( + self, + /, + arg1: onp.ToComplexND, + shape: ToShape2D | None, + dtype: _ToDType[_SCT], + copy: bool = False, + blocksize: tuple[int, int] | None = None, + ) -> None: ... + @overload # dtype: (keyword) def __init__( self, /, - arg1: Untyped, - shape: tuple[op.CanIndex, op.CanIndex] | None = None, - dtype: Untyped | None = None, + arg1: onp.ToComplexND, + shape: ToShape2D | None = None, + *, + dtype: _ToDType[_SCT], copy: bool = False, blocksize: tuple[int, int] | None = None, ) -> None: ... -class bsr_array(_bsr_base, sparray): ... -class bsr_matrix(spmatrix, _bsr_base): ... +class bsr_array(_bsr_base[_SCT], sparray, Generic[_SCT]): ... +class bsr_matrix(spmatrix[_SCT], _bsr_base[_SCT], Generic[_SCT]): ... -def isspmatrix_bsr(x: Untyped) -> bool: ... +def isspmatrix_bsr(x: object) -> TypeIs[bsr_matrix]: ... diff --git a/scipy-stubs/sparse/_compressed.pyi b/scipy-stubs/sparse/_compressed.pyi index 42fd19c3..74f10602 100644 --- a/scipy-stubs/sparse/_compressed.pyi +++ b/scipy-stubs/sparse/_compressed.pyi @@ -1,20 +1,38 @@ import abc -from typing import Literal -from typing_extensions import override +from collections.abc import Sequence +from typing import Any, Generic, Literal, TypeAlias, overload +from typing_extensions import Self, TypeVar, override -import optype as op -from scipy._typing import Untyped -from ._coo import _coo_base +import numpy as np +import optype.numpy as onp +import optype.typing as opt +from ._base import _spbase from ._data import _data_matrix, _minmax_mixin from ._index import IndexMixin +from ._typing import Int, Scalar, ToShape2D __all__: list[str] = [] -# TODO(jorenham): generic dtype -class _cs_matrix(_data_matrix, _minmax_mixin, IndexMixin): - data: Untyped - indices: Untyped - indptr: Untyped +_T = TypeVar("_T") +_SCT = TypeVar("_SCT", bound=Scalar) +_SCT_co = TypeVar("_SCT_co", bound=Scalar, default=Scalar, covariant=True) + +_ToDType: TypeAlias = type[_SCT] | np.dtype[_SCT] | onp.HasDType[np.dtype[_SCT]] +_ToMatrix: TypeAlias = _spbase[_SCT] | onp.CanArrayND[_SCT] | Sequence[onp.CanArrayND[_SCT]] | _ToMatrixPy[_SCT] +_ToMatrixPy: TypeAlias = Sequence[_T] | Sequence[Sequence[_T]] + +_ToData2B: TypeAlias = tuple[onp.ArrayND[_SCT], onp.ArrayND[Int]] # bsr +_ToData2C: TypeAlias = tuple[onp.ArrayND[_SCT], tuple[onp.ArrayND[Int], onp.ArrayND[Int]]] # csc, csr +_ToData2: TypeAlias = _ToData2B[_SCT] | _ToData2C[_SCT] +_ToData3: TypeAlias = tuple[onp.ArrayND[_SCT], onp.ArrayND[Int], onp.ArrayND[Int]] +_ToData: TypeAlias = _ToData2[_SCT] | _ToData3[_SCT] + +### + +class _cs_matrix(_data_matrix[_SCT_co], _minmax_mixin[_SCT_co], IndexMixin[_SCT_co], Generic[_SCT_co]): + data: onp.Array[Any, _SCT_co] # the `Any` shape is needed for `numpy<2.1` + indices: onp.Array1D[np.int32] + indptr: onp.Array1D[np.int32] @property @override @@ -25,20 +43,85 @@ class _cs_matrix(_data_matrix, _minmax_mixin, IndexMixin): @property def has_canonical_format(self, /) -> bool: ... @has_canonical_format.setter - def has_canonical_format(self, /, val: bool) -> None: ... + def has_canonical_format(self, val: bool, /) -> None: ... # @property def has_sorted_indices(self, /) -> bool: ... @has_sorted_indices.setter - def has_sorted_indices(self, /, val: bool) -> None: ... + def has_sorted_indices(self, val: bool, /) -> None: ... # + @overload # matrix-like (known dtype), dtype: None def __init__( self, /, - arg1: Untyped, - shape: tuple[op.CanIndex, op.CanIndex] | None = None, - dtype: Untyped | None = None, + arg1: _ToMatrix[_SCT_co] | _ToData[_SCT_co], + shape: ToShape2D | None = None, + dtype: None = None, + copy: bool = False, + ) -> None: ... + @overload # 2-d shape-like, dtype: None + def __init__( + self: _cs_matrix[np.float64], + /, + arg1: ToShape2D, + shape: None = None, + dtype: None = None, + copy: bool = False, + ) -> None: ... + @overload # matrix-like builtins.bool, dtype: type[bool] | None + def __init__( + self: _cs_matrix[np.bool_], + /, + arg1: _ToMatrixPy[bool], + shape: ToShape2D | None = None, + dtype: onp.AnyBoolDType | None = None, + copy: bool = False, + ) -> None: ... + @overload # matrix-like builtins.int, dtype: type[int] | None + def __init__( + self: _cs_matrix[np.int_], + /, + arg1: _ToMatrixPy[opt.JustInt], + shape: ToShape2D | None = None, + dtype: type[opt.JustInt] | onp.AnyIntPDType | None = None, + copy: bool = False, + ) -> None: ... + @overload # matrix-like builtins.float, dtype: type[float] | None + def __init__( + self: _cs_matrix[np.float64], + /, + arg1: _ToMatrixPy[opt.Just[float]], + shape: ToShape2D | None = None, + dtype: type[opt.Just[float]] | onp.AnyFloat64DType | None = None, + copy: bool = False, + ) -> None: ... + @overload # matrix-like builtins.complex, dtype: type[complex] | None + def __init__( + self: _cs_matrix[np.complex128], + /, + arg1: _ToMatrixPy[opt.Just[complex]], + shape: ToShape2D | None = None, + dtype: type[opt.Just[complex]] | onp.AnyComplex128DType | None = None, + copy: bool = False, + ) -> None: ... + @overload # dtype: (positional) + def __init__( + self, + /, + arg1: onp.ToComplexND, + shape: ToShape2D | None, + dtype: _ToDType[_SCT_co], + copy: bool = False, + ) -> None: ... + @overload # dtype: (keyword) + def __init__( + self, + /, + arg1: onp.ToComplexND, + shape: ToShape2D | None = None, + *, + dtype: _ToDType[_SCT_co], copy: bool = False, ) -> None: ... @@ -47,15 +130,11 @@ class _cs_matrix(_data_matrix, _minmax_mixin, IndexMixin): def count_nonzero(self, /, axis: None = None) -> int: ... # - def check_format(self, /, full_check: bool = True) -> Untyped: ... - def eliminate_zeros(self, /) -> Untyped: ... - def sum_duplicates(self, /) -> Untyped: ... - def sorted_indices(self, /) -> Untyped: ... - - # + def sorted_indices(self, /) -> Self: ... def sort_indices(self, /) -> None: ... - def prune(self, /) -> None: ... # - @override - def tocoo(self, /, copy: bool = True) -> _coo_base: ... + def check_format(self, /, full_check: bool = True) -> None: ... + def eliminate_zeros(self, /) -> None: ... + def sum_duplicates(self, /) -> None: ... + def prune(self, /) -> None: ... diff --git a/scipy-stubs/sparse/_coo.pyi b/scipy-stubs/sparse/_coo.pyi index 127b0119..36bf1734 100644 --- a/scipy-stubs/sparse/_coo.pyi +++ b/scipy-stubs/sparse/_coo.pyi @@ -1,17 +1,32 @@ -from typing import Literal -from typing_extensions import override +from collections.abc import Sequence +from typing import Any, Generic, Literal, TypeAlias, overload +from typing_extensions import TypeIs, TypeVar, override -from scipy._typing import Untyped -from ._base import sparray +import numpy as np +import optype as op +import optype.numpy as onp +import optype.typing as opt +from ._base import _spbase, sparray from ._data import _data_matrix, _minmax_mixin from ._matrix import spmatrix +from ._typing import Int, Scalar, ToShape2D __all__ = ["coo_array", "coo_matrix", "isspmatrix_coo"] -# TODO(jorenham): generic dtype and shape -class _coo_base(_data_matrix, _minmax_mixin): - coords: Untyped - data: Untyped +_T = TypeVar("_T") +_SCT = TypeVar("_SCT", bound=Scalar, default=Any) + +_ToDType: TypeAlias = type[_SCT] | np.dtype[_SCT] | onp.HasDType[np.dtype[_SCT]] +_ToMatrix: TypeAlias = _spbase[_SCT] | onp.CanArrayND[_SCT] | Sequence[onp.CanArrayND[_SCT]] | _ToMatrixPy[_SCT] +_ToMatrixPy: TypeAlias = Sequence[_T] | Sequence[Sequence[_T]] +_ToData: TypeAlias = tuple[onp.ArrayND[_SCT], tuple[onp.ArrayND[Int]] | tuple[onp.ArrayND[Int], onp.ArrayND[Int]]] + +### + +# TODO(jorenham): Make generic on `shape` +class _coo_base(_data_matrix[_SCT], _minmax_mixin[_SCT], Generic[_SCT]): + data: onp.Array1D[_SCT] + coords: tuple[onp.Array1D[np.int32]] | tuple[onp.Array1D[np.int32], onp.Array1D[np.int32]] has_canonical_format: bool @property @@ -19,30 +34,101 @@ class _coo_base(_data_matrix, _minmax_mixin): def format(self, /) -> Literal["coo"]: ... # @property - def row(self, /) -> Untyped: ... + def row(self, /) -> onp.Array1D[np.int32]: ... @row.setter - def row(self, /, new_row: Untyped) -> None: ... + def row(self, new_row: onp.ToInt1D, /) -> None: ... # @property - def col(self, /) -> Untyped: ... + def col(self, /) -> onp.Array1D[np.int32]: ... @col.setter - def col(self, /, new_col: Untyped) -> None: ... + def col(self, new_col: onp.ToInt1D, /) -> None: ... # + @overload # matrix-like (known dtype), dtype: None + def __init__( + self, + /, + arg1: _ToMatrix[_SCT] | _ToData[_SCT], + shape: ToShape2D | None = None, + dtype: None = None, + copy: bool = False, + ) -> None: ... + @overload # 2-d shape-like, dtype: None + def __init__( + self: _coo_base[np.float64], + /, + arg1: ToShape2D, + shape: None = None, + dtype: None = None, + copy: bool = False, + ) -> None: ... + @overload # matrix-like builtins.bool, dtype: type[bool] | None + def __init__( + self: _coo_base[np.bool_], + /, + arg1: _ToMatrixPy[bool], + shape: ToShape2D | None = None, + dtype: onp.AnyBoolDType | None = None, + copy: bool = False, + ) -> None: ... + @overload # matrix-like builtins.int, dtype: type[int] | None + def __init__( + self: _coo_base[np.int_], + /, + arg1: _ToMatrixPy[opt.JustInt], + shape: ToShape2D | None = None, + dtype: type[opt.JustInt] | onp.AnyIntPDType | None = None, + copy: bool = False, + ) -> None: ... + @overload # matrix-like builtins.float, dtype: type[float] | None + def __init__( + self: _coo_base[np.float64], + /, + arg1: _ToMatrixPy[opt.Just[float]], + shape: ToShape2D | None = None, + dtype: type[opt.Just[float]] | onp.AnyFloat64DType | None = None, + copy: bool = False, + ) -> None: ... + @overload # matrix-like builtins.complex, dtype: type[complex] | None + def __init__( + self: _coo_base[np.complex128], + /, + arg1: _ToMatrixPy[opt.Just[complex]], + shape: ToShape2D | None = None, + dtype: type[opt.Just[complex]] | onp.AnyComplex128DType | None = None, + copy: bool = False, + ) -> None: ... + @overload # dtype: (positional) + def __init__( + self, + /, + arg1: onp.ToComplexND, + shape: ToShape2D | None, + dtype: _ToDType[_SCT], + copy: bool = False, + ) -> None: ... + @overload # dtype: (keyword) def __init__( self, /, - arg1: Untyped, - shape: Untyped | None = None, - dtype: Untyped | None = None, + arg1: onp.ToComplexND, + shape: ToShape2D | None = None, + *, + dtype: _ToDType[_SCT], copy: bool = False, ) -> None: ... # def sum_duplicates(self, /) -> None: ... - def eliminate_zeros(self, /) -> Untyped: ... + def eliminate_zeros(self, /) -> None: ... + +class coo_array(_coo_base[_SCT], sparray, Generic[_SCT]): ... -class coo_array(_coo_base, sparray): ... -class coo_matrix(spmatrix, _coo_base): ... +class coo_matrix(spmatrix[_SCT], _coo_base[_SCT], Generic[_SCT]): + # NOTE: using `@override` together with `@overload` causes stubtest to crash... + @overload # type: ignore[explicit-override] + def getnnz(self, /, axis: None = None) -> int: ... + @overload + def getnnz(self, /, axis: op.CanIndex) -> onp.Array1D[np.intp]: ... -def isspmatrix_coo(x: Untyped) -> bool: ... +def isspmatrix_coo(x: object) -> TypeIs[coo_matrix]: ... diff --git a/scipy-stubs/sparse/_csc.pyi b/scipy-stubs/sparse/_csc.pyi index a1fca5e3..08378b5e 100644 --- a/scipy-stubs/sparse/_csc.pyi +++ b/scipy-stubs/sparse/_csc.pyi @@ -1,19 +1,29 @@ -from typing import Literal -from typing_extensions import override +from typing import Any, Generic, Literal +from typing_extensions import TypeIs, TypeVar, override -from scipy._typing import Untyped from ._base import sparray from ._compressed import _cs_matrix from ._matrix import spmatrix +from ._typing import Scalar __all__ = ["csc_array", "csc_matrix", "isspmatrix_csc"] -class _csc_base(_cs_matrix): +_SCT = TypeVar("_SCT", bound=Scalar, default=Any) + +### + +class _csc_base(_cs_matrix[_SCT], Generic[_SCT]): @property @override def format(self, /) -> Literal["csc"]: ... + @property + @override + def ndim(self, /) -> Literal[2]: ... + @property + @override + def shape(self, /) -> tuple[int, int]: ... -class csc_array(_csc_base, sparray): ... -class csc_matrix(spmatrix, _csc_base): ... +class csc_array(_csc_base[_SCT], sparray, Generic[_SCT]): ... +class csc_matrix(spmatrix[_SCT], _csc_base[_SCT], Generic[_SCT]): ... -def isspmatrix_csc(x: Untyped) -> bool: ... +def isspmatrix_csc(x: object) -> TypeIs[csc_matrix]: ... diff --git a/scipy-stubs/sparse/_csr.pyi b/scipy-stubs/sparse/_csr.pyi index a694e5a6..d19d8a95 100644 --- a/scipy-stubs/sparse/_csr.pyi +++ b/scipy-stubs/sparse/_csr.pyi @@ -1,25 +1,29 @@ -from typing import Literal -from typing_extensions import override +from typing import Any, Generic, Literal +from typing_extensions import TypeIs, TypeVar, override -from scipy._typing import Untyped from ._base import sparray -from ._bsr import _bsr_base from ._compressed import _cs_matrix from ._matrix import spmatrix +from ._typing import Scalar __all__ = ["csr_array", "csr_matrix", "isspmatrix_csr"] -# TODO(jorenham): generic dtype -class _csr_base(_cs_matrix): +_SCT = TypeVar("_SCT", bound=Scalar, default=Any) + +### + +class _csr_base(_cs_matrix[_SCT], Generic[_SCT]): @property @override def format(self, /) -> Literal["csr"]: ... - - # + @property + @override + def ndim(self, /) -> Literal[2]: ... + @property @override - def tobsr(self, /, blocksize: tuple[int, int] | None = None, copy: bool = True) -> _bsr_base: ... + def shape(self, /) -> tuple[int, int]: ... -class csr_array(_csr_base, sparray): ... -class csr_matrix(spmatrix, _csr_base): ... +class csr_array(_csr_base[_SCT], sparray, Generic[_SCT]): ... +class csr_matrix(spmatrix[_SCT], _csr_base[_SCT], Generic[_SCT]): ... -def isspmatrix_csr(x: Untyped) -> bool: ... +def isspmatrix_csr(x: object) -> TypeIs[csr_matrix]: ... diff --git a/scipy-stubs/sparse/_data.pyi b/scipy-stubs/sparse/_data.pyi index 5dbcde36..c941965f 100644 --- a/scipy-stubs/sparse/_data.pyi +++ b/scipy-stubs/sparse/_data.pyi @@ -3,13 +3,15 @@ from typing import Any, Generic, Literal, overload from typing_extensions import Never, Self, TypeVar, override import numpy as np +import optype as op from scipy._typing import Untyped -from ._base import _Scalar, _spbase +from ._base import _spbase from ._sputils import _ScalarLike +from ._typing import Scalar __all__: list[str] = [] -_SCT_co = TypeVar("_SCT_co", bound=_Scalar, default=_Scalar, covariant=True) +_SCT_co = TypeVar("_SCT_co", bound=Scalar, default=Scalar, covariant=True) ### @@ -109,10 +111,11 @@ class _data_matrix(_spbase[_SCT_co], Generic[_SCT_co]): @overload def rad2deg(self, /) -> Self: ... -class _minmax_mixin: - def max(self, /, axis: int | None = None, out: Untyped | None = None) -> Untyped: ... - def min(self, /, axis: int | None = None, out: Untyped | None = None) -> Untyped: ... - def nanmax(self, /, axis: int | None = None, out: Untyped | None = None) -> Untyped: ... - def nanmin(self, /, axis: int | None = None, out: Untyped | None = None) -> Untyped: ... - def argmax(self, /, axis: int | None = None, out: Untyped | None = None) -> Untyped: ... - def argmin(self, /, axis: int | None = None, out: Untyped | None = None) -> Untyped: ... +# TODO(jorenham) +class _minmax_mixin(Generic[_SCT_co]): + def max(self, /, axis: op.CanIndex | None = None, out: Untyped | None = None) -> Untyped: ... + def min(self, /, axis: op.CanIndex | None = None, out: Untyped | None = None) -> Untyped: ... + def nanmax(self, /, axis: op.CanIndex | None = None, out: Untyped | None = None) -> Untyped: ... + def nanmin(self, /, axis: op.CanIndex | None = None, out: Untyped | None = None) -> Untyped: ... + def argmax(self, /, axis: op.CanIndex | None = None, out: Untyped | None = None) -> Untyped: ... + def argmin(self, /, axis: op.CanIndex | None = None, out: Untyped | None = None) -> Untyped: ... diff --git a/scipy-stubs/sparse/_dia.pyi b/scipy-stubs/sparse/_dia.pyi index 31e52380..a0cc4f77 100644 --- a/scipy-stubs/sparse/_dia.pyi +++ b/scipy-stubs/sparse/_dia.pyi @@ -1,17 +1,30 @@ -from typing import Literal -from typing_extensions import override +from collections.abc import Sequence +from typing import Any, Generic, Literal, TypeAlias, overload +from typing_extensions import TypeIs, TypeVar, override -import optype as op +import numpy as np +import optype.numpy as onp +import optype.typing as opt from scipy._typing import Untyped -from ._base import sparray +from ._base import _spbase, sparray from ._data import _data_matrix from ._matrix import spmatrix +from ._typing import Int, Scalar, ToShape2D __all__ = ["dia_array", "dia_matrix", "isspmatrix_dia"] -# TODO(jorenham): generic dtype -class _dia_base(_data_matrix): - data: Untyped +_T = TypeVar("_T") +_SCT = TypeVar("_SCT", bound=Scalar, default=Any) + +_ToDType: TypeAlias = type[_SCT] | np.dtype[_SCT] | onp.HasDType[np.dtype[_SCT]] +_ToMatrix: TypeAlias = _spbase[_SCT] | onp.CanArrayND[_SCT] | Sequence[onp.CanArrayND[_SCT]] | _ToMatrixPy[_SCT] +_ToMatrixPy: TypeAlias = Sequence[_T] | Sequence[Sequence[_T]] +_ToData: TypeAlias = tuple[onp.ArrayND[_SCT], onp.ArrayND[Int]] + +### + +class _dia_base(_data_matrix[_SCT], Generic[_SCT]): + data: onp.Array2D[_SCT] offsets: Untyped @property @@ -25,16 +38,81 @@ class _dia_base(_data_matrix): def shape(self, /) -> tuple[int, int]: ... # + @overload # matrix-like (known dtype), dtype: None + def __init__( + self, + /, + arg1: _ToMatrix[_SCT] | _ToData[_SCT], + shape: ToShape2D | None = None, + dtype: None = None, + copy: bool = False, + ) -> None: ... + @overload # 2-d shape-like, dtype: None + def __init__( + self: _dia_base[np.float64], + /, + arg1: ToShape2D, + shape: None = None, + dtype: None = None, + copy: bool = False, + ) -> None: ... + @overload # matrix-like builtins.bool, dtype: type[bool] | None + def __init__( + self: _dia_base[np.bool_], + /, + arg1: _ToMatrixPy[bool], + shape: ToShape2D | None = None, + dtype: onp.AnyBoolDType | None = None, + copy: bool = False, + ) -> None: ... + @overload # matrix-like builtins.int, dtype: type[int] | None + def __init__( + self: _dia_base[np.int_], + /, + arg1: _ToMatrixPy[opt.JustInt], + shape: ToShape2D | None = None, + dtype: type[opt.JustInt] | onp.AnyIntPDType | None = None, + copy: bool = False, + ) -> None: ... + @overload # matrix-like builtins.float, dtype: type[float] | None + def __init__( + self: _dia_base[np.float64], + /, + arg1: _ToMatrixPy[opt.Just[float]], + shape: ToShape2D | None = None, + dtype: type[opt.Just[float]] | onp.AnyFloat64DType | None = None, + copy: bool = False, + ) -> None: ... + @overload # matrix-like builtins.complex, dtype: type[complex] | None + def __init__( + self: _dia_base[np.complex128], + /, + arg1: _ToMatrixPy[opt.Just[complex]], + shape: ToShape2D | None = None, + dtype: type[opt.Just[complex]] | onp.AnyComplex128DType | None = None, + copy: bool = False, + ) -> None: ... + @overload # dtype: (positional) + def __init__( + self, + /, + arg1: onp.ToComplexND, + shape: ToShape2D | None, + dtype: _ToDType[_SCT], + copy: bool = False, + ) -> None: ... + @overload # dtype: (keyword) def __init__( self, /, - arg1: Untyped, - shape: tuple[op.CanIndex, op.CanIndex] | None = None, - dtype: Untyped | None = None, + arg1: onp.ToComplexND, + shape: ToShape2D | None = None, + *, + dtype: _ToDType[_SCT], copy: bool = False, ) -> None: ... -class dia_array(_dia_base, sparray): ... -class dia_matrix(spmatrix, _dia_base): ... +class dia_array(_dia_base[_SCT], sparray, Generic[_SCT]): ... +class dia_matrix(spmatrix[_SCT], _dia_base[_SCT], Generic[_SCT]): ... -def isspmatrix_dia(x: Untyped) -> bool: ... +def isspmatrix_dia(x: object) -> TypeIs[dia_matrix]: ... diff --git a/scipy-stubs/sparse/_dok.pyi b/scipy-stubs/sparse/_dok.pyi index bab4014c..2ae8c7d5 100644 --- a/scipy-stubs/sparse/_dok.pyi +++ b/scipy-stubs/sparse/_dok.pyi @@ -1,67 +1,146 @@ # mypy: disable-error-code="misc, override" # pyright: reportIncompatibleMethodOverride=false -from collections.abc import Iterable -from typing import Literal, NoReturn -from typing_extensions import Never, Self, override +from collections.abc import Iterable, Sequence +from typing import Any, Generic, Literal, TypeAlias, overload +from typing_extensions import Never, Self, TypeIs, TypeVar, override -import optype as op -from scipy._typing import Untyped +import numpy as np +import optype.numpy as onp +import optype.typing as opt from ._base import _spbase, sparray from ._index import IndexMixin from ._matrix import spmatrix +from ._typing import Scalar, ToShape __all__ = ["dok_array", "dok_matrix", "isspmatrix_dok"] -# TODO(jorenham): generic dtype and shape -class _dok_base(_spbase, IndexMixin, dict[tuple[int, int], Untyped]): - dtype: Untyped +_T = TypeVar("_T") +_SCT = TypeVar("_SCT", bound=Scalar, default=Any) + +_ToDType: TypeAlias = type[_SCT] | np.dtype[_SCT] | onp.HasDType[np.dtype[_SCT]] +_ToMatrix: TypeAlias = _spbase[_SCT] | onp.CanArrayND[_SCT] | Sequence[onp.CanArrayND[_SCT]] | _ToMatrixPy[_SCT] +_ToMatrixPy: TypeAlias = Sequence[_T] | Sequence[Sequence[_T]] + +_Key: TypeAlias = tuple[onp.ToJustInt] | tuple[onp.ToJustInt, onp.ToJustInt] + +### + +# TODO(jorenham): Make generic on `shape` +class _dok_base(_spbase[_SCT], IndexMixin[_SCT], dict[_Key, _SCT], Generic[_SCT]): + dtype: np.dtype[_SCT] @property @override def format(self, /) -> Literal["dok"]: ... # + @overload # matrix-like (known dtype), dtype: None + def __init__( + self, + /, + arg1: _ToMatrix[_SCT], + shape: ToShape | None = None, + dtype: None = None, + copy: bool = False, + ) -> None: ... + @overload # 2-d shape-like, dtype: None + def __init__( + self: _dok_base[np.float64], + /, + arg1: ToShape, + shape: None = None, + dtype: None = None, + copy: bool = False, + ) -> None: ... + @overload # matrix-like builtins.bool, dtype: type[bool] | None + def __init__( + self: _dok_base[np.bool_], + /, + arg1: _ToMatrixPy[bool], + shape: ToShape | None = None, + dtype: onp.AnyBoolDType | None = None, + copy: bool = False, + ) -> None: ... + @overload # matrix-like builtins.int, dtype: type[int] | None + def __init__( + self: _dok_base[np.int_], + /, + arg1: _ToMatrixPy[opt.JustInt], + shape: ToShape | None = None, + dtype: type[opt.JustInt] | onp.AnyIntPDType | None = None, + copy: bool = False, + ) -> None: ... + @overload # matrix-like builtins.float, dtype: type[float] | None + def __init__( + self: _dok_base[np.float64], + /, + arg1: _ToMatrixPy[opt.Just[float]], + shape: ToShape | None = None, + dtype: type[opt.Just[float]] | onp.AnyFloat64DType | None = None, + copy: bool = False, + ) -> None: ... + @overload # matrix-like builtins.complex, dtype: type[complex] | None + def __init__( + self: _dok_base[np.complex128], + /, + arg1: _ToMatrixPy[opt.Just[complex]], + shape: ToShape | None = None, + dtype: type[opt.Just[complex]] | onp.AnyComplex128DType | None = None, + copy: bool = False, + ) -> None: ... + @overload # dtype: (positional) def __init__( self, /, - arg1: Untyped, - shape: Untyped | None = None, - dtype: Untyped | None = None, + arg1: onp.ToComplexND, + shape: ToShape | None, + dtype: _ToDType[_SCT], + copy: bool = False, + ) -> None: ... + @overload # dtype: (keyword) + def __init__( + self, + /, + arg1: onp.ToComplexND, + shape: ToShape | None = None, + *, + dtype: _ToDType[_SCT], copy: bool = False, ) -> None: ... # @override - def __delitem__(self, key: op.CanIndex, /) -> None: ... - @override - def __or__(self, other: Never, /) -> NoReturn: ... + def __delitem__(self, key: onp.ToJustInt, /) -> None: ... + + # @override - def __ror__(self, other: Never, /) -> NoReturn: ... + def __or__(self, other: Never, /) -> Never: ... @override - def __ior__(self, other: Never, /) -> Self: ... + def __ror__(self, other: Never, /) -> Never: ... @override - def get(self, key: Untyped, /, default: float = 0.0) -> Untyped: ... - @classmethod + def __ior__(self, other: Never, /) -> Never: ... # noqa: PYI034 @override - def fromkeys(cls, iterable: Iterable[tuple[int, ...]], value: int = 1, /) -> Self: ... + def update(self, /, val: Never) -> Never: ... - # + # TODO(jorenham) @override - def update(self, /, val: Untyped) -> None: ... + def get(self, key: onp.ToJustInt | _Key, /, default: onp.ToComplex = 0.0) -> _SCT: ... @override - def setdefault(self, key: Untyped, default: Untyped | None = None, /) -> Untyped: ... + def setdefault(self, key: onp.ToJustInt | _Key, default: onp.ToComplex | None = None, /) -> _SCT: ... + @classmethod + @override + def fromkeys(cls, iterable: Iterable[_Key], value: int = 1, /) -> Self: ... # - def conjtransp(self, /) -> Untyped: ... + def conjtransp(self, /) -> Self: ... -class dok_array(_dok_base, sparray): ... +class dok_array(_dok_base[_SCT], sparray, Generic[_SCT]): ... -class dok_matrix(spmatrix, _dok_base): - @property +class dok_matrix(spmatrix[_SCT], _dok_base[_SCT], Generic[_SCT]): @override - def shape(self, /) -> tuple[int, int]: ... + def get(self, key: tuple[onp.ToJustInt, onp.ToJustInt], /, default: onp.ToComplex = 0.0) -> _SCT: ... @override - def get_shape(self, /) -> tuple[int, int]: ... + def setdefault(self, key: tuple[onp.ToJustInt, onp.ToJustInt], default: onp.ToComplex | None = None, /) -> _SCT: ... -def isspmatrix_dok(x: Untyped) -> bool: ... +def isspmatrix_dok(x: object) -> TypeIs[dok_matrix]: ... diff --git a/scipy-stubs/sparse/_index.pyi b/scipy-stubs/sparse/_index.pyi index d251f8d7..7ed97838 100644 --- a/scipy-stubs/sparse/_index.pyi +++ b/scipy-stubs/sparse/_index.pyi @@ -1,22 +1,51 @@ -from collections.abc import Sequence from types import EllipsisType -from typing import Any, TypeAlias +from typing import Any, Generic, TypeAlias, overload +from typing_extensions import Buffer, Self, TypeVar import numpy as np import optype as op import optype.numpy as onp -from scipy._typing import Untyped from ._base import _spbase +from ._typing import Complex, Float, Int, Scalar -_Mask: TypeAlias = _spbase[np.bool_] | onp.CanArrayND[np.bool_] | Sequence[Sequence[bool]] -_Indexer: TypeAlias = op.CanIndex | onp.CanArrayND[np.bool_] | Sequence[bool] | onp.ArrayND[np.intp] | slice | EllipsisType -_Key: TypeAlias = tuple[_Indexer] | tuple[_Indexer, _Indexer] | Sequence[bool] | _Mask | EllipsisType +_SCT_co = TypeVar("_SCT_co", bound=Scalar, default=Scalar, covariant=True) + +_ToInt: TypeAlias = str | op.CanInt | op.CanIndex | op.CanTrunc | Buffer +_ToFloat: TypeAlias = str | op.CanFloat | op.CanIndex | Buffer +_ToComplex: TypeAlias = str | op.CanComplex | op.CanFloat | op.CanIndex | complex + +_ToIndex1D: TypeAlias = op.CanIndex | tuple[op.CanIndex] +_ToIndex2D: TypeAlias = tuple[op.CanIndex, op.CanIndex] + +_ToSlice1D: TypeAlias = slice | EllipsisType | onp.ToJustInt2D | onp.ToBool1D +_ToSlice2D: TypeAlias = ( + slice + | EllipsisType + | tuple[_ToIndex1D | _ToSlice1D, _ToSlice1D] + | tuple[_ToSlice1D, _ToSlice1D | _ToIndex1D] + | onp.ToBool2D + | list[bool] + | list[np.bool_] + | list[int] + | _spbase[np.bool_] +) ### INT_TYPES: tuple[type[int], type[np.integer[Any]]] = ... -# TODO(jorenham): generic scalar type -class IndexMixin: - def __getitem__(self, key: _Key, /) -> Untyped: ... - def __setitem__(self, key: _Key, x: onp.ToComplex | onp.ToComplex1D | onp.ToComplex2D | _spbase, /) -> None: ... +class IndexMixin(Generic[_SCT_co]): + @overload + def __getitem__(self, ix: _ToIndex2D, /) -> _SCT_co: ... + @overload + def __getitem__(self, ixs: _ToSlice2D, /) -> Self: ... + + # + @overload + def __setitem__(self: IndexMixin[Int], key: _ToIndex2D, x: _ToInt, /) -> None: ... + @overload + def __setitem__(self: IndexMixin[Float], key: _ToIndex2D, x: _ToFloat, /) -> None: ... + @overload + def __setitem__(self: IndexMixin[Complex], key: _ToIndex2D, x: _ToComplex, /) -> None: ... + @overload + def __setitem__(self, key: _ToIndex2D | _ToSlice2D, x: object, /) -> None: ... diff --git a/scipy-stubs/sparse/_lil.pyi b/scipy-stubs/sparse/_lil.pyi index 18300ac8..003d6515 100644 --- a/scipy-stubs/sparse/_lil.pyi +++ b/scipy-stubs/sparse/_lil.pyi @@ -1,19 +1,31 @@ -from typing import Literal -from typing_extensions import override +from collections.abc import Sequence +from typing import Any, Generic, Literal, TypeAlias, overload +from typing_extensions import Self, TypeIs, TypeVar, override -import optype as op -from scipy._typing import Untyped +import numpy as np +import optype.numpy as onp +import optype.typing as opt from ._base import _spbase, sparray +from ._csr import csr_array, csr_matrix from ._index import IndexMixin from ._matrix import spmatrix +from ._typing import Scalar, ToShape2D __all__ = ["isspmatrix_lil", "lil_array", "lil_matrix"] -# TODO(jorenham): generic dtype -class _lil_base(_spbase, IndexMixin): - dtype: Untyped - rows: Untyped - data: Untyped +_T = TypeVar("_T") +_SCT = TypeVar("_SCT", bound=Scalar, default=Any) + +_ToDType: TypeAlias = type[_SCT] | np.dtype[_SCT] | onp.HasDType[np.dtype[_SCT]] +_ToMatrix: TypeAlias = _spbase[_SCT] | onp.CanArrayND[_SCT] | Sequence[onp.CanArrayND[_SCT]] | _ToMatrixPy[_SCT] +_ToMatrixPy: TypeAlias = Sequence[_T] | Sequence[Sequence[_T]] + +### + +class _lil_base(_spbase[_SCT], IndexMixin[_SCT], Generic[_SCT]): + dtype: np.dtype[_SCT] + data: onp.Array1D[np.object_] + rows: onp.Array1D[np.object_] @property @override @@ -26,12 +38,77 @@ class _lil_base(_spbase, IndexMixin): def shape(self, /) -> tuple[int, int]: ... # + @overload # matrix-like (known dtype), dtype: None + def __init__( + self, + /, + arg1: _ToMatrix[_SCT], + shape: ToShape2D | None = None, + dtype: None = None, + copy: bool = False, + ) -> None: ... + @overload # 2-d shape-like, dtype: None + def __init__( + self: _lil_base[np.float64], + /, + arg1: ToShape2D, + shape: None = None, + dtype: None = None, + copy: bool = False, + ) -> None: ... + @overload # matrix-like builtins.bool, dtype: type[bool] | None + def __init__( + self: _lil_base[np.bool_], + /, + arg1: _ToMatrixPy[bool], + shape: ToShape2D | None = None, + dtype: onp.AnyBoolDType | None = None, + copy: bool = False, + ) -> None: ... + @overload # matrix-like builtins.int, dtype: type[int] | None + def __init__( + self: _lil_base[np.int_], + /, + arg1: _ToMatrixPy[opt.JustInt], + shape: ToShape2D | None = None, + dtype: type[opt.JustInt] | onp.AnyIntPDType | None = None, + copy: bool = False, + ) -> None: ... + @overload # matrix-like builtins.float, dtype: type[float] | None + def __init__( + self: _lil_base[np.float64], + /, + arg1: _ToMatrixPy[opt.Just[float]], + shape: ToShape2D | None = None, + dtype: type[opt.Just[float]] | onp.AnyFloat64DType | None = None, + copy: bool = False, + ) -> None: ... + @overload # matrix-like builtins.complex, dtype: type[complex] | None + def __init__( + self: _lil_base[np.complex128], + /, + arg1: _ToMatrixPy[opt.Just[complex]], + shape: ToShape2D | None = None, + dtype: type[opt.Just[complex]] | onp.AnyComplex128DType | None = None, + copy: bool = False, + ) -> None: ... + @overload # dtype: (positional) def __init__( self, /, - arg1: Untyped, - shape: tuple[op.CanIndex, op.CanIndex] | None = None, - dtype: Untyped | None = None, + arg1: onp.ToComplexND, + shape: ToShape2D | None, + dtype: _ToDType[_SCT], + copy: bool = False, + ) -> None: ... + @overload # dtype: (keyword) + def __init__( + self, + /, + arg1: onp.ToComplexND, + shape: ToShape2D | None = None, + *, + dtype: _ToDType[_SCT], copy: bool = False, ) -> None: ... @@ -40,10 +117,15 @@ class _lil_base(_spbase, IndexMixin): def resize(self, /, *shape: int) -> None: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride] # - def getrowview(self, /, i: int) -> Untyped: ... - def getrow(self, /, i: int) -> Untyped: ... + def getrowview(self, /, i: int) -> Self: ... + def getrow(self, /, i: onp.ToJustInt) -> csr_array[_SCT] | csr_matrix[_SCT]: ... + +class lil_array(_lil_base[_SCT], sparray, Generic[_SCT]): + @override + def getrow(self, /, i: onp.ToJustInt) -> csr_array[_SCT]: ... -class lil_array(_lil_base, sparray): ... -class lil_matrix(spmatrix, _lil_base): ... +class lil_matrix(spmatrix[_SCT], _lil_base[_SCT], Generic[_SCT]): + @override + def getrow(self, /, i: onp.ToJustInt) -> csr_matrix[_SCT]: ... -def isspmatrix_lil(x: Untyped) -> Untyped: ... +def isspmatrix_lil(x: object) -> TypeIs[lil_matrix]: ... diff --git a/scipy-stubs/sparse/_matrix.pyi b/scipy-stubs/sparse/_matrix.pyi index 5f2ca5ba..af9122a8 100644 --- a/scipy-stubs/sparse/_matrix.pyi +++ b/scipy-stubs/sparse/_matrix.pyi @@ -1,8 +1,48 @@ +# pyright: reportUnannotatedClassAttribute=false + +# needed (once) for `numpy>=2.2.0` +# mypy: disable-error-code="overload-overlap" + +from typing import Generic, overload +from typing_extensions import Self, TypeVar + +import numpy as np +import optype.numpy as onp from scipy._typing import Untyped +from ._bsr import bsr_matrix +from ._coo import coo_matrix +from ._csc import csc_matrix +from ._csr import csr_matrix +from ._dia import dia_matrix +from ._dok import dok_matrix +from ._lil import lil_matrix +from ._typing import Scalar, SPFormat, ToShape2D + +_SCT_co = TypeVar("_SCT_co", bound=Scalar, default=Scalar, covariant=True) -class spmatrix: +### + +class spmatrix(Generic[_SCT_co]): + @property + def _bsr_container(self, /) -> bsr_matrix[_SCT_co]: ... + @property + def _coo_container(self, /) -> coo_matrix[_SCT_co]: ... + @property + def _csc_container(self, /) -> csc_matrix[_SCT_co]: ... + @property + def _csr_container(self, /) -> csr_matrix[_SCT_co]: ... + @property + def _dia_container(self, /) -> dia_matrix[_SCT_co]: ... + @property + def _dok_container(self, /) -> dok_matrix[_SCT_co]: ... + @property + def _lil_container(self, /) -> lil_matrix[_SCT_co]: ... + + # @property def shape(self, /) -> tuple[int, int]: ... + def get_shape(self, /) -> tuple[int, int]: ... + def set_shape(self, /, shape: ToShape2D) -> None: ... # def __mul__(self, other: Untyped, /) -> Untyped: ... @@ -10,12 +50,18 @@ class spmatrix: def __pow__(self, power: Untyped, /) -> Untyped: ... # - def set_shape(self, /, shape: Untyped) -> None: ... - def get_shape(self, /) -> tuple[int, ...]: ... - def asfptype(self, /) -> Untyped: ... - def getmaxprint(self, /) -> Untyped: ... - def getformat(self, /) -> Untyped: ... - def getnnz(self, /, axis: Untyped | None = None) -> Untyped: ... - def getH(self, /) -> Untyped: ... - def getcol(self, /, j: int) -> Untyped: ... - def getrow(self, /, i: int) -> Untyped: ... + def getmaxprint(self, /) -> int: ... + def getformat(self, /) -> SPFormat: ... + # NOTE: `axis` is only supported by `{coo,csc,csr,lil}_matrix` + def getnnz(self, /, axis: None = None) -> int: ... + def getH(self, /) -> Self: ... + def getcol(self, /, j: onp.ToJustInt) -> csc_matrix[_SCT_co]: ... + def getrow(self, /, i: onp.ToJustInt) -> csr_matrix[_SCT_co]: ... + + # NOTE: mypy reports a false positive for overlapping overloads + @overload + def asfptype(self: spmatrix[np.bool_ | np.int8 | np.int16 | np.uint8 | np.uint16], /) -> spmatrix[np.float32]: ... + @overload + def asfptype(self: spmatrix[np.int32 | np.int64 | np.uint32 | np.uint64], /) -> spmatrix[np.float64]: ... + @overload + def asfptype(self, /) -> Self: ... diff --git a/scipy-stubs/sparse/_typing.pyi b/scipy-stubs/sparse/_typing.pyi new file mode 100644 index 00000000..0b37a5a3 --- /dev/null +++ b/scipy-stubs/sparse/_typing.pyi @@ -0,0 +1,29 @@ +# NOTE(scipy-stubs): This ia a module only exists `if typing.TYPE_CHECKING: ...` +from typing import Any, Literal, TypeAlias +from typing_extensions import TypeAliasType + +import numpy as np +import optype as op + +__all__ = "Complex", "Float", "Int", "SPFormat", "Scalar", "ToShape", "ToShape1D", "ToShape2D" + +# NOTE: This is roughly speaking equivalent to `[u]int8 | [u]int16 | [u]int32 | [u]int64` (on most modern platforms) +Int: TypeAlias = np.integer[Any] +# NOTE: NOT equivalent to `floating`! It's considered invalid to use `float16` in `scipy.sparse`. +Float: TypeAlias = np.float32 | np.float64 | np.longdouble +# NOTE: This is (almost always) equivalent to `complex64 | complex128 | clongdouble`. +Complex: TypeAlias = np.complexfloating[Any, Any] +# NOTE: Roughly speaking, this is equivalent to `number[Any] - float16` (the `-` denotes the set difference analogue for types) +# NOTE: This (almost always) matches `._sputils.supported_dtypes` +# NOTE: The `TypeAliasType` is used to avoid long error messages. +Scalar = TypeAliasType("Scalar", np.bool_ | Int | Float | Complex) + +### + +SPFormat: TypeAlias = Literal["bsr", "coo", "csc", "csr", "dia", "dok", "lil"] + +### + +ToShape1D: TypeAlias = tuple[op.CanIndex] +ToShape2D: TypeAlias = tuple[op.CanIndex, op.CanIndex] +ToShape: TypeAlias = ToShape1D | ToShape2D