Skip to content

Commit d40f3a3

Browse files
jorenhamjorenham
committed
🐛 sparse: fix overridden dict method siagnatures in dok_{array,matrix} 
1 parent e1988a0 commit d40f3a3

File tree

1 file changed

+157
-15
lines changed

1 file changed

+157
-15
lines changed

‎scipy-stubs/sparse/_dok.pyi

Lines changed: 157 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,13 @@
1-
# mypy: disable-error-code="misc, override"
1+
# NOTE: Adding `@override` to `@overload`ed methods will crash stubtest (basedmypy 1.13.0)
2+
# mypy: disable-error-code="misc, override, explicit-override"
23
# pyright: reportIncompatibleMethodOverride=false
34

45
from collections.abc import Iterable, Sequence
56
from typing import Any, Generic, Literal, TypeAlias, overload
67
from typing_extensions import Never, Self, TypeIs, TypeVar, override
78

89
import numpy as np
10+
import optype as op
911
import optype.numpy as onp
1012
import optype.typing as opt
1113
from ._base import _spbase, sparray
@@ -15,14 +17,25 @@ from ._typing import Scalar, ShapeDOK, ToShape1dNd
1517

1618
__all__ = ["dok_array", "dok_matrix", "isspmatrix_dok"]
1719

20+
###
21+
1822
_T = TypeVar("_T")
1923
_SCT = TypeVar("_SCT", bound=Scalar, default=Any)
2024
_ShapeT_co = TypeVar("_ShapeT_co", bound=ShapeDOK, default=ShapeDOK, covariant=True)
2125

26+
_1D: TypeAlias = tuple[int] # noqa: PYI042
27+
_2D: TypeAlias = tuple[int, int] # noqa: PYI042
28+
2229
_ToDType: TypeAlias = type[_SCT] | np.dtype[_SCT] | onp.HasDType[np.dtype[_SCT]]
2330
_ToMatrix: TypeAlias = _spbase[_SCT] | onp.CanArrayND[_SCT] | Sequence[onp.CanArrayND[_SCT]] | _ToMatrixPy[_SCT]
2431
_ToMatrixPy: TypeAlias = Sequence[_T] | Sequence[Sequence[_T]]
2532

33+
_ToKey1D: TypeAlias = onp.ToJustInt | tuple[onp.ToJustInt]
34+
_ToKey2D: TypeAlias = tuple[onp.ToJustInt, onp.ToJustInt]
35+
36+
_ToKeys1D: TypeAlias = Iterable[_ToKey1D]
37+
_ToKeys2D: TypeAlias = Iterable[_ToKey2D]
38+
2639
###
2740

2841
class _dok_base(_spbase[_SCT, _ShapeT_co], IndexMixin[_SCT, _ShapeT_co], dict[ShapeDOK, _SCT], Generic[_SCT, _ShapeT_co]):
@@ -31,11 +44,9 @@ class _dok_base(_spbase[_SCT, _ShapeT_co], IndexMixin[_SCT, _ShapeT_co], dict[Sh
3144
@property
3245
@override
3346
def format(self, /) -> Literal["dok"]: ...
34-
#
3547
@property
3648
@override
3749
def ndim(self, /) -> Literal[1, 2]: ...
38-
#
3950
@property
4051
@override
4152
def shape(self, /) -> _ShapeT_co: ...
@@ -141,27 +152,158 @@ class _dok_base(_spbase[_SCT, _ShapeT_co], IndexMixin[_SCT, _ShapeT_co], dict[Sh
141152
def __ror__(self, other: Never, /) -> Never: ...
142153
@override
143154
def __ior__(self, other: Never, /) -> Never: ... # noqa: PYI034
155+
156+
#
157+
@overload
158+
def count_nonzero(self, /, axis: None = None) -> int: ...
159+
@overload
160+
def count_nonzero(self, /, axis: op.CanIndex) -> onp.Array1D[np.intp]: ...
161+
162+
#
144163
@override
145164
def update(self, /, val: Never) -> Never: ...
146165

147-
# TODO(jorenham)
148-
@override
149-
def get(self, key: onp.ToJustInt | ShapeDOK, /, default: onp.ToComplex = 0.0) -> _SCT: ...
150-
@override
151-
def setdefault(self, key: onp.ToJustInt | ShapeDOK, default: onp.ToComplex | None = None, /) -> _SCT: ...
152-
@classmethod
153-
@override
154-
def fromkeys(cls, iterable: Iterable[ShapeDOK], value: int = 1, /) -> Self: ...
166+
#
167+
@overload
168+
def setdefault(self: _dok_base[Any, _1D], key: _ToKey1D, default: None = None, /) -> _SCT | None: ...
169+
@overload
170+
def setdefault(self: _dok_base[Any, _2D], key: _ToKey2D, default: None = None, /) -> _SCT | None: ...
171+
@overload
172+
def setdefault(self, key: _ToKey1D | _ToKey2D, default: None = None, /) -> _SCT | None: ...
173+
@overload
174+
def setdefault(self: _dok_base[Any, _1D], key: _ToKey1D, default: _T, /) -> _SCT | _T: ...
175+
@overload
176+
def setdefault(self: _dok_base[Any, _2D], key: _ToKey2D, default: _T, /) -> _SCT | _T: ...
177+
@overload
178+
def setdefault(self, key: _ToKey1D | _ToKey2D, default: _T, /) -> _SCT | _T: ...
179+
180+
#
181+
@overload
182+
def get(self: _dok_base[Any, _1D], key: _ToKey1D, /, default: float = 0.0) -> _SCT | float: ...
183+
@overload
184+
def get(self: _dok_base[Any, _2D], key: _ToKey2D, /, default: float = 0.0) -> _SCT | float: ...
185+
@overload
186+
def get(self, key: _ToKey1D | _ToKey2D, /, default: float = 0.0) -> _SCT | float: ...
187+
@overload
188+
def get(self: _dok_base[Any, _1D], key: _ToKey1D, /, default: _T) -> _SCT | _T: ...
189+
@overload
190+
def get(self: _dok_base[Any, _2D], key: _ToKey2D, /, default: _T) -> _SCT | _T: ...
191+
@overload
192+
def get(self, key: _ToKey1D | _ToKey2D, /, default: _T) -> _SCT | _T: ...
155193

156194
#
157195
def conjtransp(self, /) -> Self: ...
158196

159-
class dok_array(_dok_base[_SCT, _ShapeT_co], sparray, Generic[_SCT, _ShapeT_co]): ...
197+
#
198+
@overload
199+
@classmethod
200+
def fromkeys(cls: type[_dok_base[np.bool_, _1D]], ks: _ToKeys1D, v: onp.ToBool, /) -> _dok_base[np.bool_, _1D]: ...
201+
@overload
202+
@classmethod
203+
def fromkeys(cls: type[_dok_base[np.bool_, _2D]], ks: _ToKeys2D, v: onp.ToBool, /) -> _dok_base[np.bool_, _2D]: ...
204+
@overload
205+
@classmethod
206+
def fromkeys(cls: type[_dok_base[np.int_, _1D]], ks: _ToKeys1D, v: opt.JustInt = 1, /) -> _dok_base[np.int_, _1D]: ...
207+
@overload
208+
@classmethod
209+
def fromkeys(cls: type[_dok_base[np.int_, _2D]], ks: _ToKeys2D, v: opt.JustInt = 1, /) -> _dok_base[np.int_, _2D]: ...
210+
@overload
211+
@classmethod
212+
def fromkeys(cls: type[_dok_base[np.float64, _1D]], ks: _ToKeys1D, v: opt.JustFloat, /) -> _dok_base[np.float64, _1D]: ...
213+
@overload
214+
@classmethod
215+
def fromkeys(cls: type[_dok_base[np.float64, _2D]], ks: _ToKeys2D, v: opt.JustFloat, /) -> _dok_base[np.float64, _2D]: ...
216+
@overload
217+
@classmethod
218+
def fromkeys(
219+
cls: type[_dok_base[np.complex128, _1D]],
220+
ks: _ToKeys1D,
221+
v: opt.JustComplex,
222+
/,
223+
) -> _dok_base[np.complex128, _1D]: ...
224+
@overload
225+
@classmethod
226+
def fromkeys(
227+
cls: type[_dok_base[np.complex128, _2D]],
228+
ks: _ToKeys2D,
229+
v: opt.JustComplex,
230+
/,
231+
) -> _dok_base[np.complex128, _2D]: ...
232+
@overload
233+
@classmethod
234+
def fromkeys(cls: type[_dok_base[_SCT, _1D]], ks: _ToKeys1D, v: _SCT, /) -> _dok_base[_SCT, _1D]: ...
235+
@overload
236+
@classmethod
237+
def fromkeys(cls: type[_dok_base[_SCT, _2D]], ks: _ToKeys2D, v: _SCT, /) -> _dok_base[_SCT, _2D]: ...
160238

161-
class dok_matrix(_dok_base[_SCT, tuple[int, int]], spmatrix[_SCT], Generic[_SCT]):
239+
#
240+
class dok_array(_dok_base[_SCT, _ShapeT_co], sparray, Generic[_SCT, _ShapeT_co]):
241+
# NOTE: This horrible code duplication is required due to the lack of higher-kinded typing (HKT) support.
242+
# https://github.com/python/typing/issues/548
243+
@overload
244+
@classmethod
245+
def fromkeys(cls: type[dok_array[np.bool_, _1D]], ks: _ToKeys1D, v: onp.ToBool, /) -> dok_array[np.bool_, _1D]: ...
246+
@overload
247+
@classmethod
248+
def fromkeys(cls: type[dok_array[np.bool_, _2D]], ks: _ToKeys2D, v: onp.ToBool, /) -> dok_array[np.bool_, _2D]: ...
249+
@overload
250+
@classmethod
251+
def fromkeys(cls: type[dok_array[np.int_, _1D]], ks: _ToKeys1D, v: opt.JustInt = 1, /) -> dok_array[np.int_, _1D]: ...
252+
@overload
253+
@classmethod
254+
def fromkeys(cls: type[dok_array[np.int_, _2D]], ks: _ToKeys2D, v: opt.JustInt = 1, /) -> dok_array[np.int_, _2D]: ...
255+
@overload
256+
@classmethod
257+
def fromkeys(cls: type[dok_array[np.float64, _1D]], ks: _ToKeys1D, v: opt.JustFloat, /) -> dok_array[np.float64, _1D]: ...
258+
@overload
259+
@classmethod
260+
def fromkeys(cls: type[dok_array[np.float64, _2D]], ks: _ToKeys2D, v: opt.JustFloat, /) -> dok_array[np.float64, _2D]: ...
261+
@overload
262+
@classmethod
263+
def fromkeys(
264+
cls: type[dok_array[np.complex128, _1D]],
265+
ks: _ToKeys1D,
266+
v: opt.JustComplex,
267+
/,
268+
) -> dok_array[np.complex128, _1D]: ...
269+
@overload
270+
@classmethod
271+
def fromkeys(
272+
cls: type[dok_array[np.complex128, _2D]],
273+
ks: _ToKeys2D,
274+
v: opt.JustComplex,
275+
/,
276+
) -> dok_array[np.complex128, _2D]: ...
277+
@overload
278+
@classmethod
279+
def fromkeys(cls: type[dok_array[_SCT, _1D]], ks: _ToKeys1D, v: _SCT, /) -> dok_array[_SCT, _1D]: ...
280+
@overload
281+
@classmethod
282+
def fromkeys(cls: type[dok_array[_SCT, _2D]], ks: _ToKeys2D, v: _SCT, /) -> dok_array[_SCT, _2D]: ...
283+
284+
#
285+
class dok_matrix(_dok_base[_SCT, _2D], spmatrix[_SCT], Generic[_SCT]):
162286
@override
163-
def get(self, key: tuple[onp.ToJustInt, onp.ToJustInt], /, default: onp.ToComplex = 0.0) -> _SCT: ...
287+
def get(self, key: _ToKey2D, /, default: onp.ToComplex = 0.0) -> _SCT: ...
164288
@override
165-
def setdefault(self, key: tuple[onp.ToJustInt, onp.ToJustInt], default: onp.ToComplex | None = None, /) -> _SCT: ...
289+
def setdefault(self, key: _ToKey2D, default: onp.ToComplex | None = None, /) -> _SCT: ...
290+
291+
#
292+
@overload
293+
@classmethod
294+
def fromkeys(cls: type[dok_matrix[np.bool_]], ks: _ToKeys2D, v: onp.ToBool, /) -> dok_matrix[np.bool_]: ...
295+
@overload
296+
@classmethod
297+
def fromkeys(cls: type[dok_matrix[np.int_]], ks: _ToKeys2D, v: opt.JustInt = 1, /) -> dok_matrix[np.int_]: ...
298+
@overload
299+
@classmethod
300+
def fromkeys(cls: type[dok_matrix[np.float64]], ks: _ToKeys2D, v: opt.JustFloat, /) -> dok_matrix[np.float64]: ...
301+
@overload
302+
@classmethod
303+
def fromkeys(cls: type[dok_matrix[np.complex128]], ks: _ToKeys2D, v: opt.JustComplex, /) -> dok_matrix[np.complex128]: ...
304+
@overload
305+
@classmethod
306+
def fromkeys(cls, ks: _ToKeys2D, v: _SCT, /) -> Self: ...
166307

308+
#
167309
def isspmatrix_dok(x: object) -> TypeIs[dok_matrix]: ...

0 commit comments

Comments
 (0)