Skip to content

Commit 4a89cbe

Browse files
jorenhamjorenham
authored
🏷️ fix stubtest errors in numpy.lib._twodim_base_impl (#245)
1 parent 96d3db5 commit 4a89cbe

File tree

3 files changed

+91
-65
lines changed

3 files changed

+91
-65
lines changed

‎.mypyignore-todo

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8,8 +8,6 @@ numpy(\..+)?\.floating.as_integer_ratio
88
numpy(\..+)?\.complexfloating.__hash__
99
numpy(\..+)?\.complexfloating.__complex__
1010

11-
numpy(\.lib\._twodim_base_impl)?\.tri(l|u)
12-
1311
numpy\._?core(\._multiarray_umath|\.multiarray)\.error
1412
numpy(\._core(\.memmap)?|\.matlib)?\.memmap\.__new__
1513
numpy(\.matrixlib(\.defmatrix)?|\.matlib)?\.matrix\.__new__
@@ -94,7 +92,6 @@ numpy.matlib.hanning
9492
numpy.matlib.in1d
9593
numpy.matlib.split
9694
numpy.matlib.stack
97-
numpy.matlib.tri(l|u)
9895

9996
numpy.matrixlib.defmatrix.mat
10097

‎src/numpy-stubs/lib/_twodim_base_impl.pyi

Lines changed: 89 additions & 60 deletions
Original file line numberDiff line numberDiff line change
@@ -1,30 +1,31 @@
11
from collections.abc import Callable, Sequence
2-
from typing import Any, Literal as L, TypeAlias, overload
2+
from typing import Literal as L, TypeAlias, overload
33
from typing_extensions import TypeVar
44

55
import numpy as np
66
from _numtype import (
7+
Array,
78
CoComplex_1d,
89
CoFloat64_1d,
910
Is,
1011
ToBool_1d,
12+
ToBool_1nd,
13+
ToBytes_1nd,
14+
ToComplex128_1nd,
1115
ToComplex_1d,
16+
ToFloat64_1nd,
1217
ToFloat64_2d,
1318
ToFloating_1d,
1419
ToIntP_1d,
20+
ToIntP_1nd,
1521
ToInteger_1d,
1622
ToObject_1d,
23+
ToStr_1nd,
1724
_ToArray1_1d,
25+
_ToArray1_1nd,
1826
)
1927
from numpy import _OrderCF # noqa: ICN003
20-
from numpy._typing import (
21-
ArrayLike,
22-
DTypeLike,
23-
NDArray,
24-
_ArrayLike,
25-
_DTypeLike,
26-
_SupportsArrayFunc as _CanArrayFunc,
27-
)
28+
from numpy._typing import ArrayLike, DTypeLike, _ArrayLike, _DTypeLike, _SupportsArrayFunc as _CanArrayFunc
2829

2930
__all__ = [
3031
"diag",
@@ -52,8 +53,8 @@ _NumericT = TypeVar("_NumericT", bound=_CoComplex)
5253

5354
# The returned arrays dtype must be compatible with `np.equal`
5455
_MaskFunc: TypeAlias = Callable[
55-
[NDArray[np.intp], _T],
56-
NDArray[np.number | np.bool | np.timedelta64 | np.datetime64 | np.object_],
56+
[Array[np.intp], _T],
57+
Array[np.number | np.bool | np.timedelta64 | np.datetime64 | np.object_],
5758
]
5859

5960
_CoInt: TypeAlias = np.integer | np.bool
@@ -63,15 +64,15 @@ _CoComplex: TypeAlias = np.number | np.bool
6364
###
6465

6566
@overload
66-
def fliplr(m: _ArrayLike[_ScalarT]) -> NDArray[_ScalarT]: ...
67+
def fliplr(m: _ArrayLike[_ScalarT]) -> Array[_ScalarT]: ...
6768
@overload
68-
def fliplr(m: ArrayLike) -> NDArray[Any]: ...
69+
def fliplr(m: ArrayLike) -> Array: ...
6970

7071
#
7172
@overload
72-
def flipud(m: _ArrayLike[_ScalarT]) -> NDArray[_ScalarT]: ...
73+
def flipud(m: _ArrayLike[_ScalarT]) -> Array[_ScalarT]: ...
7374
@overload
74-
def flipud(m: ArrayLike) -> NDArray[Any]: ...
75+
def flipud(m: ArrayLike) -> Array: ...
7576

7677
#
7778
@overload
@@ -84,7 +85,7 @@ def eye(
8485
*,
8586
device: L["cpu"] | None = None,
8687
like: _CanArrayFunc | None = None,
87-
) -> NDArray[np.float64]: ...
88+
) -> Array[np.float64]: ...
8889
@overload
8990
def eye(
9091
N: int,
@@ -95,7 +96,7 @@ def eye(
9596
*,
9697
device: L["cpu"] | None = None,
9798
like: _CanArrayFunc | None = None,
98-
) -> NDArray[_ScalarT]: ...
99+
) -> Array[_ScalarT]: ...
99100
@overload
100101
def eye(
101102
N: int,
@@ -106,7 +107,7 @@ def eye(
106107
order: _OrderCF = "C",
107108
device: L["cpu"] | None = None,
108109
like: _CanArrayFunc | None = None,
109-
) -> NDArray[_ScalarT]: ...
110+
) -> Array[_ScalarT]: ...
110111
@overload
111112
def eye(
112113
N: int,
@@ -117,19 +118,19 @@ def eye(
117118
*,
118119
device: L["cpu"] | None = None,
119120
like: _CanArrayFunc | None = None,
120-
) -> NDArray[Any]: ...
121+
) -> Array: ...
121122

122123
#
123124
@overload
124-
def diag(v: _ArrayLike[_ScalarT], k: int = 0) -> NDArray[_ScalarT]: ...
125+
def diag(v: _ArrayLike[_ScalarT], k: int = 0) -> Array[_ScalarT]: ...
125126
@overload
126-
def diag(v: ArrayLike, k: int = 0) -> NDArray[Any]: ...
127+
def diag(v: ArrayLike, k: int = 0) -> Array: ...
127128

128129
#
129130
@overload
130-
def diagflat(v: _ArrayLike[_ScalarT], k: int = 0) -> NDArray[_ScalarT]: ...
131+
def diagflat(v: _ArrayLike[_ScalarT], k: int = 0) -> Array[_ScalarT]: ...
131132
@overload
132-
def diagflat(v: ArrayLike, k: int = 0) -> NDArray[Any]: ...
133+
def diagflat(v: ArrayLike, k: int = 0) -> Array: ...
133134

134135
#
135136
@overload
@@ -140,7 +141,7 @@ def tri(
140141
dtype: type[Is[float]] | None = ...,
141142
*,
142143
like: _CanArrayFunc | None = None,
143-
) -> NDArray[np.float64]: ...
144+
) -> Array[np.float64]: ...
144145
@overload
145146
def tri(
146147
N: int,
@@ -149,7 +150,7 @@ def tri(
149150
dtype: _DTypeLike[_ScalarT],
150151
*,
151152
like: _CanArrayFunc | None = None,
152-
) -> NDArray[_ScalarT]: ...
153+
) -> Array[_ScalarT]: ...
153154
@overload
154155
def tri(
155156
N: int,
@@ -158,7 +159,7 @@ def tri(
158159
*,
159160
dtype: _DTypeLike[_ScalarT],
160161
like: _CanArrayFunc | None = None,
161-
) -> NDArray[_ScalarT]: ...
162+
) -> Array[_ScalarT]: ...
162163
@overload
163164
def tri(
164165
N: int,
@@ -167,31 +168,57 @@ def tri(
167168
dtype: DTypeLike = ...,
168169
*,
169170
like: _CanArrayFunc | None = None,
170-
) -> NDArray[Any]: ...
171+
) -> Array: ...
171172

172173
#
173174
@overload
174-
def tril(v: _ArrayLike[_ScalarT], k: int = 0) -> NDArray[_ScalarT]: ...
175+
def tril(m: _ToArray1_1nd[_ScalarT], k: int = 0) -> Array[_ScalarT]: ...
176+
@overload
177+
def tril(m: ToBool_1nd, k: int = 0) -> Array[np.bool]: ...
178+
@overload
179+
def tril(m: ToIntP_1nd, k: int = 0) -> Array[np.intp]: ...
180+
@overload
181+
def tril(m: ToFloat64_1nd, k: int = 0) -> Array[np.float64]: ...
182+
@overload
183+
def tril(m: ToComplex128_1nd, k: int = 0) -> Array[np.complex128]: ...
175184
@overload
176-
def tril(v: ArrayLike, k: int = 0) -> NDArray[Any]: ...
185+
def tril(m: ToBytes_1nd, k: int = 0) -> Array[np.bytes_]: ...
186+
@overload
187+
def tril(m: ToStr_1nd, k: int = 0) -> Array[np.str_]: ...
188+
@overload
189+
def tril(m: ArrayLike, k: int = 0) -> Array: ...
177190

178191
#
179192
@overload
180-
def triu(v: _ArrayLike[_ScalarT], k: int = 0) -> NDArray[_ScalarT]: ...
193+
def triu(m: _ToArray1_1nd[_ScalarT], k: int = 0) -> Array[_ScalarT]: ...
194+
@overload
195+
def triu(m: ToBool_1nd, k: int = 0) -> Array[np.bool]: ...
196+
@overload
197+
def triu(m: ToIntP_1nd, k: int = 0) -> Array[np.intp]: ...
181198
@overload
182-
def triu(v: ArrayLike, k: int = 0) -> NDArray[Any]: ...
199+
def triu(m: ToFloat64_1nd, k: int = 0) -> Array[np.float64]: ...
200+
@overload
201+
def triu(m: ToComplex128_1nd, k: int = 0) -> Array[np.complex128]: ...
202+
@overload
203+
def triu(m: ToBytes_1nd, k: int = 0) -> Array[np.bytes_]: ...
204+
@overload
205+
def triu(m: ToStr_1nd, k: int = 0) -> Array[np.str_]: ...
206+
@overload
207+
def triu(m: ArrayLike, k: int = 0) -> Array: ...
183208

184209
#
185210
@overload
186-
def vander(x: ToBool_1d, N: int | None = None, increasing: bool = False) -> NDArray[np.intp]: ...
211+
def vander(x: ToBool_1d, N: int | None = None, increasing: bool = False) -> Array[np.intp]: ...
187212
@overload
188-
def vander(x: ToInteger_1d, N: int | None = None, increasing: bool = False) -> NDArray[np.signedinteger]: ...
213+
def vander(x: ToInteger_1d, N: int | None = None, increasing: bool = False) -> Array[np.signedinteger]: ...
189214
@overload
190-
def vander(x: ToFloating_1d, N: int | None = None, increasing: bool = False) -> NDArray[np.floating]: ...
215+
def vander(x: ToFloating_1d, N: int | None = None, increasing: bool = False) -> Array[np.floating]: ...
191216
@overload
192-
def vander(x: ToComplex_1d, N: int | None = None, increasing: bool = False) -> NDArray[np.complexfloating]: ...
217+
def vander(x: ToComplex_1d, N: int | None = None, increasing: bool = False) -> Array[np.complexfloating]: ...
193218
@overload
194-
def vander(x: ToObject_1d, N: int | None = None, increasing: bool = False) -> NDArray[np.object_]: ...
219+
def vander(x: ToObject_1d, N: int | None = None, increasing: bool = False) -> Array[np.object_]: ...
220+
221+
#
195222
@overload
196223
def histogram2d(
197224
x: _ToArray1_1d[_ComplexT],
@@ -200,7 +227,7 @@ def histogram2d(
200227
range: ToFloat64_2d | None = None,
201228
density: bool | None = None,
202229
weights: CoFloat64_1d | None = None,
203-
) -> tuple[NDArray[np.float64], NDArray[_ComplexT], NDArray[_ComplexT]]: ...
230+
) -> tuple[Array[np.float64], Array[_ComplexT], Array[_ComplexT]]: ...
204231
@overload
205232
def histogram2d(
206233
x: _ToArray1_1d[_ComplexT | _CoFloat],
@@ -209,7 +236,7 @@ def histogram2d(
209236
range: ToFloat64_2d | None = None,
210237
density: bool | None = None,
211238
weights: CoFloat64_1d | None = None,
212-
) -> tuple[NDArray[np.float64], NDArray[_ComplexT], NDArray[_ComplexT]]: ...
239+
) -> tuple[Array[np.float64], Array[_ComplexT], Array[_ComplexT]]: ...
213240
@overload
214241
def histogram2d(
215242
x: _ToArray1_1d[_InexactT],
@@ -218,7 +245,7 @@ def histogram2d(
218245
range: ToFloat64_2d | None = None,
219246
density: bool | None = None,
220247
weights: CoFloat64_1d | None = None,
221-
) -> tuple[NDArray[np.float64], NDArray[_InexactT], NDArray[_InexactT]]: ...
248+
) -> tuple[Array[np.float64], Array[_InexactT], Array[_InexactT]]: ...
222249
@overload
223250
def histogram2d(
224251
x: _ToArray1_1d[_InexactT | _CoInt],
@@ -227,7 +254,7 @@ def histogram2d(
227254
range: ToFloat64_2d | None = None,
228255
density: bool | None = None,
229256
weights: CoFloat64_1d | None = None,
230-
) -> tuple[NDArray[np.float64], NDArray[_InexactT], NDArray[_InexactT]]: ...
257+
) -> tuple[Array[np.float64], Array[_InexactT], Array[_InexactT]]: ...
231258
@overload
232259
def histogram2d(
233260
x: CoFloat64_1d,
@@ -236,7 +263,7 @@ def histogram2d(
236263
range: ToFloat64_2d | None = None,
237264
density: bool | None = None,
238265
weights: CoFloat64_1d | None = None,
239-
) -> tuple[NDArray[np.float64], NDArray[np.float64], NDArray[np.float64]]: ...
266+
) -> tuple[Array[np.float64], Array[np.float64], Array[np.float64]]: ...
240267
@overload
241268
def histogram2d(
242269
x: Sequence[complex],
@@ -245,7 +272,7 @@ def histogram2d(
245272
range: ToFloat64_2d | None = None,
246273
density: bool | None = None,
247274
weights: CoFloat64_1d | None = None,
248-
) -> tuple[NDArray[np.float64], NDArray[np.complex128 | np.float64], NDArray[np.complex128 | np.float64]]: ...
275+
) -> tuple[Array[np.float64], Array[np.complex128 | np.float64], Array[np.complex128 | np.float64]]: ...
249276
@overload
250277
def histogram2d(
251278
x: CoComplex_1d,
@@ -254,7 +281,7 @@ def histogram2d(
254281
range: ToFloat64_2d | None = None,
255282
density: bool | None = None,
256283
weights: CoFloat64_1d | None = None,
257-
) -> tuple[NDArray[np.float64], NDArray[_NumericT], NDArray[_NumericT]]: ...
284+
) -> tuple[Array[np.float64], Array[_NumericT], Array[_NumericT]]: ...
258285
@overload
259286
def histogram2d(
260287
x: _ToArray1_1d[_InexactT],
@@ -263,7 +290,7 @@ def histogram2d(
263290
range: ToFloat64_2d | None = None,
264291
density: bool | None = None,
265292
weights: CoFloat64_1d | None = None,
266-
) -> tuple[NDArray[np.float64], NDArray[_NumericT | _InexactT], NDArray[_NumericT | _InexactT]]: ...
293+
) -> tuple[Array[np.float64], Array[_NumericT | _InexactT], Array[_NumericT | _InexactT]]: ...
267294
@overload
268295
def histogram2d(
269296
x: ToIntP_1d | Sequence[float],
@@ -272,7 +299,7 @@ def histogram2d(
272299
range: ToFloat64_2d | None = None,
273300
density: bool | None = None,
274301
weights: CoFloat64_1d | None = None,
275-
) -> tuple[NDArray[np.float64], NDArray[_NumericT | np.float64], NDArray[_NumericT | np.float64]]: ...
302+
) -> tuple[Array[np.float64], Array[_NumericT | np.float64], Array[_NumericT | np.float64]]: ...
276303
@overload
277304
def histogram2d(
278305
x: Sequence[complex],
@@ -282,9 +309,9 @@ def histogram2d(
282309
density: bool | None = None,
283310
weights: CoFloat64_1d | None = None,
284311
) -> tuple[
285-
NDArray[np.float64],
286-
NDArray[_NumericT | np.complex128 | np.float64],
287-
NDArray[_NumericT | np.complex128 | np.float64],
312+
Array[np.float64],
313+
Array[_NumericT | np.complex128 | np.float64],
314+
Array[_NumericT | np.complex128 | np.float64],
288315
]: ...
289316
@overload
290317
def histogram2d(
@@ -294,7 +321,7 @@ def histogram2d(
294321
range: ToFloat64_2d | None = None,
295322
density: bool | None = None,
296323
weights: CoFloat64_1d | None = None,
297-
) -> tuple[NDArray[np.float64], NDArray[np.bool], NDArray[np.bool]]: ...
324+
) -> tuple[Array[np.float64], Array[np.bool], Array[np.bool]]: ...
298325
@overload
299326
def histogram2d(
300327
x: CoComplex_1d,
@@ -303,7 +330,7 @@ def histogram2d(
303330
range: ToFloat64_2d | None = None,
304331
density: bool | None = None,
305332
weights: CoFloat64_1d | None = None,
306-
) -> tuple[NDArray[np.float64], NDArray[np.intp | np.bool], NDArray[np.intp | np.bool]]: ...
333+
) -> tuple[Array[np.float64], Array[np.intp | np.bool], Array[np.intp | np.bool]]: ...
307334
@overload
308335
def histogram2d(
309336
x: CoComplex_1d,
@@ -312,7 +339,7 @@ def histogram2d(
312339
range: ToFloat64_2d | None = None,
313340
density: bool | None = None,
314341
weights: CoFloat64_1d | None = None,
315-
) -> tuple[NDArray[np.float64], NDArray[np.float64 | np.intp | np.bool], NDArray[np.float64 | np.intp | np.bool]]: ...
342+
) -> tuple[Array[np.float64], Array[np.float64 | np.intp | np.bool], Array[np.float64 | np.intp | np.bool]]: ...
316343
@overload
317344
def histogram2d(
318345
x: CoComplex_1d,
@@ -322,21 +349,23 @@ def histogram2d(
322349
density: bool | None = None,
323350
weights: CoFloat64_1d | None = None,
324351
) -> tuple[
325-
NDArray[np.float64],
326-
NDArray[np.complex128 | np.float64 | np.intp | np.bool],
327-
NDArray[np.complex128 | np.float64 | np.intp | np.bool],
352+
Array[np.float64],
353+
Array[np.complex128 | np.float64 | np.intp | np.bool],
354+
Array[np.complex128 | np.float64 | np.intp | np.bool],
328355
]: ...
329356

330357
# NOTE: we're assuming/demanding here the `mask_func` returns
331358
# an ndarray of shape `(n, n)`; otherwise there is the possibility
332359
# of the output tuple having more or less than 2 elements
333360
@overload
334-
def mask_indices(n: int, mask_func: _MaskFunc[int], k: int = 0) -> tuple[NDArray[np.intp], NDArray[np.intp]]: ...
361+
def mask_indices(n: int, mask_func: _MaskFunc[int], k: int = 0) -> tuple[Array[np.intp], Array[np.intp]]: ...
335362
@overload
336-
def mask_indices(n: int, mask_func: _MaskFunc[_T], k: _T) -> tuple[NDArray[np.intp], NDArray[np.intp]]: ...
363+
def mask_indices(n: int, mask_func: _MaskFunc[_T], k: _T) -> tuple[Array[np.intp], Array[np.intp]]: ...
364+
365+
#
366+
def tril_indices(n: int, k: int = 0, m: int | None = None) -> tuple[Array[np.intp], Array[np.intp]]: ...
367+
def triu_indices(n: int, k: int = 0, m: int | None = None) -> tuple[Array[np.intp], Array[np.intp]]: ...
337368

338369
#
339-
def tril_indices(n: int, k: int = 0, m: int | None = None) -> tuple[NDArray[np.intp], NDArray[np.intp]]: ...
340-
def triu_indices(n: int, k: int = 0, m: int | None = None) -> tuple[NDArray[np.intp], NDArray[np.intp]]: ...
341-
def tril_indices_from(arr: NDArray[Any], k: int = 0) -> tuple[NDArray[np.intp], NDArray[np.intp]]: ...
342-
def triu_indices_from(arr: NDArray[Any], k: int = 0) -> tuple[NDArray[np.intp], NDArray[np.intp]]: ...
370+
def tril_indices_from(arr: Array, k: int = 0) -> tuple[Array[np.intp], Array[np.intp]]: ...
371+
def triu_indices_from(arr: Array, k: int = 0) -> tuple[Array[np.intp], Array[np.intp]]: ...

0 commit comments

Comments
 (0)