Skip to content

Commit 1a6cb7b

Browse files
jorenhamjorenham
committed
🎨 np._core.einsumfunc cleanup
1 parent cd48b3d commit 1a6cb7b

File tree

1 file changed

+66
-65
lines changed

1 file changed

+66
-65
lines changed

‎src/numpy-stubs/_core/einsumfunc.pyi

Lines changed: 66 additions & 65 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,11 @@
1+
from _typeshed import Incomplete
12
from collections.abc import Sequence
23
from typing import Any, Literal, Literal as L, TypeAlias, overload
34
from typing_extensions import TypeVar
45

5-
import numpy as np
6+
import _numtype as _nt
7+
from numpy import _OrderKACF # noqa: ICN003
68
from numpy._typing import (
7-
NDArray,
89
_ArrayLikeBool_co,
910
_ArrayLikeComplex_co,
1011
_ArrayLikeFloat_co,
@@ -22,9 +23,9 @@ from numpy._typing import (
2223

2324
__all__ = ["einsum", "einsum_path"]
2425

25-
_ArrayT = TypeVar("_ArrayT", bound=NDArray[np.bool | np.number])
26+
_ArrayT = TypeVar("_ArrayT", bound=_nt.Array[_nt.co_complex])
2627

27-
_OptimizeKind: TypeAlias = bool | Literal["greedy", "optimal"] | Sequence[Any] | None
28+
_OptimizeKind: TypeAlias = bool | Literal["greedy", "optimal"] | Sequence[Incomplete]
2829
_CastingSafe: TypeAlias = Literal["no", "equiv", "safe", "same_kind"]
2930
_CastingUnsafe: TypeAlias = Literal["unsafe"]
3031

@@ -39,77 +40,77 @@ def einsum(
3940
subscripts: str | _ArrayLikeInt_co,
4041
/,
4142
*operands: _ArrayLikeBool_co,
42-
out: None = ...,
43-
dtype: _DTypeLikeBool | None = ...,
44-
order: np._OrderKACF = ...,
45-
casting: _CastingSafe = ...,
46-
optimize: _OptimizeKind = ...,
47-
) -> Any: ...
43+
out: None = None,
44+
optimize: _OptimizeKind | None = None,
45+
dtype: _DTypeLikeBool | None = None,
46+
order: _OrderKACF = "K",
47+
casting: _CastingSafe = "safe",
48+
) -> Incomplete: ...
4849
@overload
4950
def einsum(
5051
subscripts: str | _ArrayLikeInt_co,
5152
/,
5253
*operands: _ArrayLikeUInt_co,
53-
out: None = ...,
54-
dtype: _DTypeLikeUInt | None = ...,
55-
order: np._OrderKACF = ...,
56-
casting: _CastingSafe = ...,
57-
optimize: _OptimizeKind = ...,
58-
) -> Any: ...
54+
out: None = None,
55+
dtype: _DTypeLikeUInt | None = None,
56+
order: _OrderKACF = "K",
57+
casting: _CastingSafe = "safe",
58+
optimize: _OptimizeKind | None = None,
59+
) -> Incomplete: ...
5960
@overload
6061
def einsum(
6162
subscripts: str | _ArrayLikeInt_co,
6263
/,
6364
*operands: _ArrayLikeInt_co,
64-
out: None = ...,
65-
dtype: _DTypeLikeInt | None = ...,
66-
order: np._OrderKACF = ...,
67-
casting: _CastingSafe = ...,
68-
optimize: _OptimizeKind = ...,
69-
) -> Any: ...
65+
out: None = None,
66+
dtype: _DTypeLikeInt | None = None,
67+
order: _OrderKACF = "K",
68+
casting: _CastingSafe = "safe",
69+
optimize: _OptimizeKind | None = None,
70+
) -> Incomplete: ...
7071
@overload
7172
def einsum(
7273
subscripts: str | _ArrayLikeInt_co,
7374
/,
7475
*operands: _ArrayLikeFloat_co,
75-
out: None = ...,
76-
dtype: _DTypeLikeFloat | None = ...,
77-
order: np._OrderKACF = ...,
78-
casting: _CastingSafe = ...,
79-
optimize: _OptimizeKind = ...,
80-
) -> Any: ...
76+
out: None = None,
77+
dtype: _DTypeLikeFloat | None = None,
78+
order: _OrderKACF = "K",
79+
casting: _CastingSafe = "safe",
80+
optimize: _OptimizeKind | None = None,
81+
) -> Incomplete: ...
8182
@overload
8283
def einsum(
8384
subscripts: str | _ArrayLikeInt_co,
8485
/,
8586
*operands: _ArrayLikeComplex_co,
86-
out: None = ...,
87-
dtype: _DTypeLikeComplex | None = ...,
88-
order: np._OrderKACF = ...,
89-
casting: _CastingSafe = ...,
90-
optimize: _OptimizeKind = ...,
91-
) -> Any: ...
87+
out: None = None,
88+
dtype: _DTypeLikeComplex | None = None,
89+
order: _OrderKACF = "K",
90+
casting: _CastingSafe = "safe",
91+
optimize: _OptimizeKind | None = None,
92+
) -> Incomplete: ...
9293
@overload
9394
def einsum(
9495
subscripts: str | _ArrayLikeInt_co,
9596
/,
9697
*operands: Any,
9798
casting: _CastingUnsafe,
98-
dtype: _DTypeLikeComplex_co | None = ...,
99-
out: None = ...,
100-
order: np._OrderKACF = ...,
101-
optimize: _OptimizeKind = ...,
102-
) -> Any: ...
99+
dtype: _DTypeLikeComplex_co | None = None,
100+
out: None = None,
101+
order: _OrderKACF = "K",
102+
optimize: _OptimizeKind | None = None,
103+
) -> Incomplete: ...
103104
@overload
104105
def einsum(
105106
subscripts: str | _ArrayLikeInt_co,
106107
/,
107108
*operands: _ArrayLikeComplex_co,
108109
out: _ArrayT,
109-
dtype: _DTypeLikeComplex_co | None = ...,
110-
order: np._OrderKACF = ...,
111-
casting: _CastingSafe = ...,
112-
optimize: _OptimizeKind = ...,
110+
dtype: _DTypeLikeComplex_co | None = None,
111+
order: _OrderKACF = "K",
112+
casting: _CastingSafe = "safe",
113+
optimize: _OptimizeKind | None = None,
113114
) -> _ArrayT: ...
114115
@overload
115116
def einsum(
@@ -118,42 +119,42 @@ def einsum(
118119
*operands: Any,
119120
out: _ArrayT,
120121
casting: _CastingUnsafe,
121-
dtype: _DTypeLikeComplex_co | None = ...,
122-
order: np._OrderKACF = ...,
123-
optimize: _OptimizeKind = ...,
122+
dtype: _DTypeLikeComplex_co | None = None,
123+
order: _OrderKACF = "K",
124+
optimize: _OptimizeKind | None = None,
124125
) -> _ArrayT: ...
125126
@overload
126127
def einsum(
127128
subscripts: str | _ArrayLikeInt_co,
128129
/,
129130
*operands: _ArrayLikeObject_co,
130-
out: None = ...,
131-
dtype: _DTypeLikeObject | None = ...,
132-
order: np._OrderKACF = ...,
133-
casting: _CastingSafe = ...,
134-
optimize: _OptimizeKind = ...,
135-
) -> Any: ...
131+
out: None = None,
132+
dtype: _DTypeLikeObject | None = None,
133+
order: _OrderKACF = "K",
134+
casting: _CastingSafe = "safe",
135+
optimize: _OptimizeKind | None = None,
136+
) -> Incomplete: ...
136137
@overload
137138
def einsum(
138139
subscripts: str | _ArrayLikeInt_co,
139140
/,
140141
*operands: Any,
141142
casting: _CastingUnsafe,
142-
dtype: _DTypeLikeObject | None = ...,
143-
out: None = ...,
144-
order: np._OrderKACF = ...,
145-
optimize: _OptimizeKind = ...,
146-
) -> Any: ...
143+
dtype: _DTypeLikeObject | None = None,
144+
out: None = None,
145+
order: _OrderKACF = "K",
146+
optimize: _OptimizeKind | None = None,
147+
) -> Incomplete: ...
147148
@overload
148149
def einsum(
149150
subscripts: str | _ArrayLikeInt_co,
150151
/,
151152
*operands: _ArrayLikeObject_co,
152153
out: _ArrayT,
153-
dtype: _DTypeLikeObject | None = ...,
154-
order: np._OrderKACF = ...,
155-
casting: _CastingSafe = ...,
156-
optimize: _OptimizeKind = ...,
154+
dtype: _DTypeLikeObject | None = None,
155+
order: _OrderKACF = "K",
156+
casting: _CastingSafe = "safe",
157+
optimize: _OptimizeKind | None = None,
157158
) -> _ArrayT: ...
158159
@overload
159160
def einsum(
@@ -162,9 +163,9 @@ def einsum(
162163
*operands: Any,
163164
out: _ArrayT,
164165
casting: _CastingUnsafe,
165-
dtype: _DTypeLikeObject | None = ...,
166-
order: np._OrderKACF = ...,
167-
optimize: _OptimizeKind = ...,
166+
dtype: _DTypeLikeObject | None = None,
167+
order: _OrderKACF = "K",
168+
optimize: _OptimizeKind | None = None,
168169
) -> _ArrayT: ...
169170

170171
# NOTE: `einsum_call` is a hidden kwarg unavailable for public use.
@@ -177,4 +178,4 @@ def einsum_path(
177178
*operands: _ArrayLikeComplex_co | _DTypeLikeObject,
178179
optimize: _OptimizeKind = "greedy",
179180
einsum_call: L[False] = False,
180-
) -> tuple[list[Any], str]: ...
181+
) -> tuple[list[Incomplete], str]: ...

0 commit comments

Comments
 (0)