Skip to content

Commit 7521a86

Browse files
authored
Merge pull request #523 from numpy/_numtype-direct-imports
2 parents 2f849d4 + 95dc843 commit 7521a86

33 files changed

Lines changed: 3738 additions & 4252 deletions

pyproject.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -225,6 +225,7 @@ preview = true
225225
"sys",
226226
"numpy",
227227
"numpy.typing",
228+
"_numtype",
228229
]
229230
[tool.ruff.lint.flake8-import-conventions.extend-aliases]
230231
"ctypes" = "ct"
Lines changed: 46 additions & 47 deletions
Original file line numberDiff line numberDiff line change
@@ -1,97 +1,96 @@
11
from typing import Any
22
from typing_extensions import assert_type
33

4+
import _numtype as _nt
45
import numpy as np
5-
import numpy.typing as npt
6-
from _numtype import Array1D
76
from numpy.lib._arraysetops_impl import UniqueAllResult, UniqueCountsResult, UniqueInverseResult
87

9-
AR_b: npt.NDArray[np.bool]
10-
AR_i_: npt.NDArray[np.intp]
11-
AR_f8: npt.NDArray[np.float64]
12-
AR_M: npt.NDArray[np.datetime64]
13-
AR_O: npt.NDArray[np.object_]
8+
AR_b: _nt.Array[np.bool]
9+
AR_i_: _nt.Array[np.intp]
10+
AR_f8: _nt.Array[np.float64]
11+
AR_M: _nt.Array[np.datetime64]
12+
AR_O: _nt.Array[np.object_]
1413
AR_LIKE_i_: list[int]
1514

1615
###
1716

18-
assert_type(np.intersect1d(AR_i_, AR_i_), Array1D[np.intp])
19-
assert_type(np.intersect1d(AR_f8, AR_i_), Array1D[np.float64])
20-
assert_type(np.intersect1d(AR_M, AR_M, assume_unique=True), Array1D[np.datetime64])
17+
assert_type(np.intersect1d(AR_i_, AR_i_), _nt.Array1D[np.intp])
18+
assert_type(np.intersect1d(AR_f8, AR_i_), _nt.Array1D[np.float64])
19+
assert_type(np.intersect1d(AR_M, AR_M, assume_unique=True), _nt.Array1D[np.datetime64])
2120
assert_type(
2221
np.intersect1d(AR_f8, AR_f8, return_indices=True),
23-
tuple[Array1D[np.float64], Array1D[np.intp], Array1D[np.intp]],
22+
tuple[_nt.Array1D[np.float64], _nt.Array1D[np.intp], _nt.Array1D[np.intp]],
2423
)
2524

26-
assert_type(np.union1d(AR_i_, AR_i_), Array1D[np.intp])
27-
assert_type(np.union1d(AR_f8, AR_i_), Array1D[np.float64])
28-
assert_type(np.union1d(AR_M, AR_M), Array1D[np.datetime64])
25+
assert_type(np.union1d(AR_i_, AR_i_), _nt.Array1D[np.intp])
26+
assert_type(np.union1d(AR_f8, AR_i_), _nt.Array1D[np.float64])
27+
assert_type(np.union1d(AR_M, AR_M), _nt.Array1D[np.datetime64])
2928

30-
assert_type(np.ediff1d(AR_b), Array1D[np.int8])
31-
assert_type(np.ediff1d(AR_M), Array1D[np.timedelta64])
32-
assert_type(np.ediff1d(AR_O), Array1D[np.object_])
33-
assert_type(np.ediff1d(AR_i_, to_end=[1, 2, 3]), Array1D[np.intp])
34-
assert_type(np.ediff1d(AR_LIKE_i_, to_begin=[0, 1]), Array1D[np.intp])
29+
assert_type(np.ediff1d(AR_b), _nt.Array1D[np.int8])
30+
assert_type(np.ediff1d(AR_M), _nt.Array1D[np.timedelta64])
31+
assert_type(np.ediff1d(AR_O), _nt.Array1D[np.object_])
32+
assert_type(np.ediff1d(AR_i_, to_end=[1, 2, 3]), _nt.Array1D[np.intp])
33+
assert_type(np.ediff1d(AR_LIKE_i_, to_begin=[0, 1]), _nt.Array1D[np.intp])
3534

36-
assert_type(np.setxor1d(AR_i_, AR_i_), Array1D[np.intp])
37-
assert_type(np.setxor1d(AR_f8, AR_i_), Array1D[np.float64])
38-
assert_type(np.setxor1d(AR_M, AR_M, assume_unique=True), Array1D[np.datetime64])
35+
assert_type(np.setxor1d(AR_i_, AR_i_), _nt.Array1D[np.intp])
36+
assert_type(np.setxor1d(AR_f8, AR_i_), _nt.Array1D[np.float64])
37+
assert_type(np.setxor1d(AR_M, AR_M, assume_unique=True), _nt.Array1D[np.datetime64])
3938

40-
assert_type(np.setdiff1d(AR_i_, AR_i_), Array1D[np.intp])
41-
assert_type(np.setdiff1d(AR_f8, AR_i_), Array1D[np.float64])
42-
assert_type(np.setdiff1d(AR_M, AR_M, assume_unique=True), Array1D[np.datetime64])
39+
assert_type(np.setdiff1d(AR_i_, AR_i_), _nt.Array1D[np.intp])
40+
assert_type(np.setdiff1d(AR_f8, AR_i_), _nt.Array1D[np.float64])
41+
assert_type(np.setdiff1d(AR_M, AR_M, assume_unique=True), _nt.Array1D[np.datetime64])
4342

44-
assert_type(np.isin(AR_i_, AR_i_), npt.NDArray[np.bool])
45-
assert_type(np.isin(AR_f8, AR_i_), npt.NDArray[np.bool])
46-
assert_type(np.isin(AR_M, AR_M, assume_unique=True), npt.NDArray[np.bool])
47-
assert_type(np.isin(AR_f8, AR_LIKE_i_, invert=True), npt.NDArray[np.bool])
43+
assert_type(np.isin(AR_i_, AR_i_), _nt.Array[np.bool])
44+
assert_type(np.isin(AR_f8, AR_i_), _nt.Array[np.bool])
45+
assert_type(np.isin(AR_M, AR_M, assume_unique=True), _nt.Array[np.bool])
46+
assert_type(np.isin(AR_f8, AR_LIKE_i_, invert=True), _nt.Array[np.bool])
4847

49-
assert_type(np.unique(AR_f8), npt.NDArray[np.float64])
50-
assert_type(np.unique(AR_f8, return_index=True), tuple[npt.NDArray[np.float64], npt.NDArray[np.intp]])
51-
assert_type(np.unique(AR_f8, return_inverse=True), tuple[npt.NDArray[np.float64], npt.NDArray[np.intp]])
52-
assert_type(np.unique(AR_f8, return_counts=True), tuple[npt.NDArray[np.float64], npt.NDArray[np.intp]])
53-
assert_type(np.unique(AR_LIKE_i_, axis=0), npt.NDArray[Any])
54-
assert_type(np.unique(AR_LIKE_i_, return_index=True), tuple[npt.NDArray[Any], npt.NDArray[np.intp]])
55-
assert_type(np.unique(AR_LIKE_i_, return_inverse=True), tuple[npt.NDArray[Any], npt.NDArray[np.intp]])
56-
assert_type(np.unique(AR_LIKE_i_, return_counts=True), tuple[npt.NDArray[Any], npt.NDArray[np.intp]])
48+
assert_type(np.unique(AR_f8), _nt.Array[np.float64])
49+
assert_type(np.unique(AR_f8, return_index=True), tuple[_nt.Array[np.float64], _nt.Array[np.intp]])
50+
assert_type(np.unique(AR_f8, return_inverse=True), tuple[_nt.Array[np.float64], _nt.Array[np.intp]])
51+
assert_type(np.unique(AR_f8, return_counts=True), tuple[_nt.Array[np.float64], _nt.Array[np.intp]])
52+
assert_type(np.unique(AR_LIKE_i_, axis=0), _nt.Array[Any])
53+
assert_type(np.unique(AR_LIKE_i_, return_index=True), tuple[_nt.Array[Any], _nt.Array[np.intp]])
54+
assert_type(np.unique(AR_LIKE_i_, return_inverse=True), tuple[_nt.Array[Any], _nt.Array[np.intp]])
55+
assert_type(np.unique(AR_LIKE_i_, return_counts=True), tuple[_nt.Array[Any], _nt.Array[np.intp]])
5756
assert_type(
5857
np.unique(AR_f8, return_index=True, return_inverse=True),
59-
tuple[npt.NDArray[np.float64], npt.NDArray[np.intp], npt.NDArray[np.intp]],
58+
tuple[_nt.Array[np.float64], _nt.Array[np.intp], _nt.Array[np.intp]],
6059
)
6160
assert_type(
6261
np.unique(AR_LIKE_i_, return_index=True, return_inverse=True),
63-
tuple[npt.NDArray[Any], npt.NDArray[np.intp], npt.NDArray[np.intp]],
62+
tuple[_nt.Array[Any], _nt.Array[np.intp], _nt.Array[np.intp]],
6463
)
6564
assert_type(
6665
np.unique(AR_f8, return_index=True, return_counts=True),
67-
tuple[npt.NDArray[np.float64], npt.NDArray[np.intp], npt.NDArray[np.intp]],
66+
tuple[_nt.Array[np.float64], _nt.Array[np.intp], _nt.Array[np.intp]],
6867
)
6968
assert_type(
7069
np.unique(AR_LIKE_i_, return_index=True, return_counts=True),
71-
tuple[npt.NDArray[Any], npt.NDArray[np.intp], npt.NDArray[np.intp]],
70+
tuple[_nt.Array[Any], _nt.Array[np.intp], _nt.Array[np.intp]],
7271
)
7372
assert_type(
7473
np.unique(AR_f8, return_inverse=True, return_counts=True),
75-
tuple[npt.NDArray[np.float64], npt.NDArray[np.intp], npt.NDArray[np.intp]],
74+
tuple[_nt.Array[np.float64], _nt.Array[np.intp], _nt.Array[np.intp]],
7675
)
7776
assert_type(
7877
np.unique(AR_LIKE_i_, return_inverse=True, return_counts=True),
79-
tuple[npt.NDArray[Any], npt.NDArray[np.intp], npt.NDArray[np.intp]],
78+
tuple[_nt.Array[Any], _nt.Array[np.intp], _nt.Array[np.intp]],
8079
)
8180
assert_type(
8281
np.unique(AR_f8, return_index=True, return_inverse=True, return_counts=True),
83-
tuple[npt.NDArray[np.float64], npt.NDArray[np.intp], npt.NDArray[np.intp], npt.NDArray[np.intp]],
82+
tuple[_nt.Array[np.float64], _nt.Array[np.intp], _nt.Array[np.intp], _nt.Array[np.intp]],
8483
)
8584
assert_type(
8685
np.unique(AR_LIKE_i_, return_index=True, return_inverse=True, return_counts=True),
87-
tuple[npt.NDArray[Any], npt.NDArray[np.intp], npt.NDArray[np.intp], npt.NDArray[np.intp]],
86+
tuple[_nt.Array[Any], _nt.Array[np.intp], _nt.Array[np.intp], _nt.Array[np.intp]],
8887
)
8988

9089
assert_type(np.unique_all(AR_f8), UniqueAllResult[np.float64])
9190
assert_type(np.unique_counts(AR_f8), UniqueCountsResult[np.float64])
9291
assert_type(np.unique_inverse(AR_f8), UniqueInverseResult[np.float64])
93-
assert_type(np.unique_values(AR_f8), Array1D[np.float64])
92+
assert_type(np.unique_values(AR_f8), _nt.Array1D[np.float64])
9493
assert_type(np.unique_all(AR_LIKE_i_), UniqueAllResult[np.intp])
9594
assert_type(np.unique_counts(AR_LIKE_i_), UniqueCountsResult[np.intp])
9695
assert_type(np.unique_inverse(AR_LIKE_i_), UniqueInverseResult[np.intp])
97-
assert_type(np.unique_values(AR_LIKE_i_), Array1D[np.intp])
96+
assert_type(np.unique_values(AR_LIKE_i_), _nt.Array1D[np.intp])

src/numpy-stubs/@test/static/accept/linalg.pyi

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,14 @@
11
from typing import Any, Literal, TypeAlias
22
from typing_extensions import TypeVar, assert_type
33

4+
import _numtype as _nt
45
import numpy as np
56
import numpy.typing as npt
6-
from _numtype import Array, AtLeast2D
77
from numpy.linalg._linalg import EigResult, EighResult, QRResult, SVDResult, SlogdetResult
88

99
_ScalarT = TypeVar("_ScalarT", bound=np.generic)
10-
_Array_2d: TypeAlias = Array[_ScalarT, tuple[int, int]]
11-
_Array_2nd: TypeAlias = Array[_ScalarT, AtLeast2D]
10+
_Array_2d: TypeAlias = _nt.Array[_ScalarT, tuple[int, int]]
11+
_Array_2nd: TypeAlias = _nt.Array[_ScalarT, _nt.AtLeast2D]
1212

1313
###
1414

src/numpy-stubs/@test/static/accept/polynomial_polybase.pyi

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -122,7 +122,7 @@ assert_type(PS_lag.integ(SC_i_co, SC_f_co), npp.Laguerre)
122122
assert_type(PS_poly.deriv(), npp.Polynomial)
123123
assert_type(PS_herm.deriv(SC_i_co), npp.Hermite)
124124

125-
assert_type(PS_poly.roots(), _Array1D[np.float64] | _Array1D[np.complex128])
125+
assert_type(PS_poly.roots(), _Array1D[np.float64 | np.complex128])
126126

127127
assert_type(
128128
PS_poly.linspace(),

src/numpy-stubs/@test/static/reject/lib_function_base.pyi

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ np.piecewise(AR_f8, AR_b_list, [fn_ar_i], 42, _=None) # pyright: ignore[reportC
2828

2929
np.interp(AR_f8, AR_c16, AR_f8) # type: ignore[arg-type] # pyright: ignore[reportArgumentType, reportCallIssue]
3030
np.interp(AR_c16, AR_f8, AR_f8) # type: ignore[arg-type] # pyright: ignore[reportArgumentType, reportCallIssue]
31-
np.interp(AR_f8, AR_f8, AR_f8, period=AR_c16) # type: ignore[call-overload] # pyright: ignore[reportArgumentType, reportCallIssue]
31+
np.interp(AR_f8, AR_f8, AR_f8, period=AR_c16) # type: ignore[arg-type] # pyright: ignore[reportArgumentType, reportCallIssue]
3232
np.interp(AR_f8, AR_f8, AR_O) # type: ignore[arg-type] # pyright: ignore[reportArgumentType, reportCallIssue]
3333

3434
np.average(AR_m) # type: ignore[arg-type] # pyright: ignore[reportArgumentType, reportCallIssue]

src/numpy-stubs/_core/_asarray.pyi

Lines changed: 24 additions & 68 deletions
Original file line numberDiff line numberDiff line change
@@ -2,15 +2,15 @@ from collections.abc import Iterable
22
from typing import Any, Final, Literal as L, TypeAlias, overload
33
from typing_extensions import TypeVar
44

5+
import _numtype as _nt
56
import numpy as np
6-
from _numtype import ToBool_nd, ToBytes_nd, ToComplex128_nd, ToFloat64_nd, ToInt_nd, ToObject_nd, ToStr_nd
7-
from numpy._typing import DTypeLike, NDArray, _ArrayLike, _DTypeLike, _SupportsArrayFunc as _Like
7+
from numpy._typing import _SupportsArrayFunc as _Like
88

99
__all__ = ["require"]
1010

1111
###
1212

13-
_ArrayT = TypeVar("_ArrayT", bound=NDArray[Any])
13+
_ArrayT = TypeVar("_ArrayT", bound=_nt.Array[Any])
1414
_ScalarT = TypeVar("_ScalarT", bound=np.generic[Any])
1515

1616
_Req: TypeAlias = L[
@@ -30,89 +30,45 @@ POSSIBLE_FLAGS: Final[dict[_ReqE, L["C", "F", "A", "W", "O", "E"]]]
3030

3131
@overload
3232
def require(
33-
a: _ArrayT,
34-
dtype: None = None,
35-
requirements: _ToReqs | None = None,
36-
*,
37-
like: _Like | None = None,
33+
a: _ArrayT, dtype: None = None, requirements: _ToReqs | None = None, *, like: _Like | None = None
3834
) -> _ArrayT: ...
3935
@overload
4036
def require(
41-
a: ToBool_nd,
42-
dtype: None = None,
43-
requirements: _ToReqsE | None = None,
44-
*,
45-
like: _Like | None = None,
46-
) -> NDArray[np.bool]: ...
37+
a: _nt.ToBool_nd, dtype: None = None, requirements: _ToReqsE | None = None, *, like: _Like | None = None
38+
) -> _nt.Array[np.bool]: ...
4739
@overload
4840
def require(
49-
a: ToInt_nd,
50-
dtype: None = None,
51-
requirements: _ToReqsE | None = None,
52-
*,
53-
like: _Like | None = None,
54-
) -> NDArray[np.intp]: ...
41+
a: _nt.ToInt_nd, dtype: None = None, requirements: _ToReqsE | None = None, *, like: _Like | None = None
42+
) -> _nt.Array[np.intp]: ...
5543
@overload
5644
def require(
57-
a: ToFloat64_nd,
58-
dtype: None = None,
59-
requirements: _ToReqsE | None = None,
60-
*,
61-
like: _Like | None = None,
62-
) -> NDArray[np.float64]: ...
45+
a: _nt.ToFloat64_nd, dtype: None = None, requirements: _ToReqsE | None = None, *, like: _Like | None = None
46+
) -> _nt.Array[np.float64]: ...
6347
@overload
6448
def require(
65-
a: ToComplex128_nd,
66-
dtype: None = None,
67-
requirements: _ToReqsE | None = None,
68-
*,
69-
like: _Like | None = None,
70-
) -> NDArray[np.complex128]: ...
49+
a: _nt.ToComplex128_nd, dtype: None = None, requirements: _ToReqsE | None = None, *, like: _Like | None = None
50+
) -> _nt.Array[np.complex128]: ...
7151
@overload
7252
def require(
73-
a: ToBytes_nd,
74-
dtype: None = None,
75-
requirements: _ToReqsE | None = None,
76-
*,
77-
like: _Like | None = None,
78-
) -> NDArray[np.bytes_]: ...
53+
a: _nt.ToBytes_nd, dtype: None = None, requirements: _ToReqsE | None = None, *, like: _Like | None = None
54+
) -> _nt.Array[np.bytes_]: ...
7955
@overload
8056
def require(
81-
a: ToStr_nd,
82-
dtype: None = None,
83-
requirements: _ToReqsE | None = None,
84-
*,
85-
like: _Like | None = None,
86-
) -> NDArray[np.str_]: ...
57+
a: _nt.ToStr_nd, dtype: None = None, requirements: _ToReqsE | None = None, *, like: _Like | None = None
58+
) -> _nt.Array[np.str_]: ...
8759
@overload
8860
def require(
89-
a: ToObject_nd,
90-
dtype: None = None,
91-
requirements: _ToReqsE | None = None,
92-
*,
93-
like: _Like | None = None,
94-
) -> NDArray[np.object_]: ...
61+
a: _nt.ToObject_nd, dtype: None = None, requirements: _ToReqsE | None = None, *, like: _Like | None = None
62+
) -> _nt.Array[np.object_]: ...
9563
@overload
9664
def require(
97-
a: _ArrayLike[_ScalarT],
98-
dtype: None = None,
99-
requirements: _ToReqsE | None = None,
100-
*,
101-
like: _Like | None = None,
102-
) -> NDArray[_ScalarT]: ...
65+
a: _nt._ToArray_nd[_ScalarT], dtype: None = None, requirements: _ToReqsE | None = None, *, like: _Like | None = None
66+
) -> _nt.Array[_ScalarT]: ...
10367
@overload
10468
def require(
105-
a: object,
106-
dtype: _DTypeLike[_ScalarT],
107-
requirements: _ToReqsE | None = None,
108-
*,
109-
like: _Like | None = None,
110-
) -> NDArray[_ScalarT]: ...
69+
a: object, dtype: _nt._ToDType[_ScalarT], requirements: _ToReqsE | None = None, *, like: _Like | None = None
70+
) -> _nt.Array[_ScalarT]: ...
11171
@overload
11272
def require(
113-
a: object,
114-
dtype: DTypeLike | None = None,
115-
requirements: _ToReqsE | None = None,
116-
*,
117-
like: _Like | None = None,
118-
) -> NDArray[Any]: ...
73+
a: object, dtype: _nt.ToDType | None = None, requirements: _ToReqsE | None = None, *, like: _Like | None = None
74+
) -> _nt.Array[Any]: ...

0 commit comments

Comments
 (0)