codekingpro/portable-devtools

1from collections.abc import Sequence
2from typing import Any, Literal, TypeAlias, TypeVar, overload
3
4import numpy as np
5from numpy import _OrderKACF, number
6from numpy._typing import (
7    NDArray,
8    _ArrayLikeBool_co,
9    _ArrayLikeComplex_co,
10    _ArrayLikeFloat_co,
11    _ArrayLikeInt_co,
12    _ArrayLikeObject_co,
13    _ArrayLikeUInt_co,
14    _DTypeLikeBool,
15    _DTypeLikeComplex,
16    _DTypeLikeComplex_co,
17    _DTypeLikeFloat,
18    _DTypeLikeInt,
19    _DTypeLikeObject,
20    _DTypeLikeUInt,
21)
22
23__all__ = ["einsum", "einsum_path"]
24
25_ArrayT = TypeVar(
26    "_ArrayT",
27    bound=NDArray[np.bool | number],
28)
29
30_OptimizeKind: TypeAlias = bool | Literal["greedy", "optimal"] | Sequence[Any] | None
31_CastingSafe: TypeAlias = Literal["no", "equiv", "safe", "same_kind"]
32_CastingUnsafe: TypeAlias = Literal["unsafe"]
33
34# TODO: Properly handle the `casting`-based combinatorics
35# TODO: We need to evaluate the content `__subscripts` in order
36# to identify whether or an array or scalar is returned. At a cursory
37# glance this seems like something that can quite easily be done with
38# a mypy plugin.
39# Something like `is_scalar = bool(__subscripts.partition("->")[-1])`
40@overload
41def einsum(
42    subscripts: str | _ArrayLikeInt_co,
43    /,
44    *operands: _ArrayLikeBool_co,
45    out: None = None,
46    dtype: _DTypeLikeBool | None = ...,
47    order: _OrderKACF = ...,
48    casting: _CastingSafe = ...,
49    optimize: _OptimizeKind = False,
50) -> Any: ...
51@overload
52def einsum(
53    subscripts: str | _ArrayLikeInt_co,
54    /,
55    *operands: _ArrayLikeUInt_co,
56    out: None = None,
57    dtype: _DTypeLikeUInt | None = ...,
58    order: _OrderKACF = ...,
59    casting: _CastingSafe = ...,
60    optimize: _OptimizeKind = False,
61) -> Any: ...
62@overload
63def einsum(
64    subscripts: str | _ArrayLikeInt_co,
65    /,
66    *operands: _ArrayLikeInt_co,
67    out: None = None,
68    dtype: _DTypeLikeInt | None = ...,
69    order: _OrderKACF = ...,
70    casting: _CastingSafe = ...,
71    optimize: _OptimizeKind = False,
72) -> Any: ...
73@overload
74def einsum(
75    subscripts: str | _ArrayLikeInt_co,
76    /,
77    *operands: _ArrayLikeFloat_co,
78    out: None = None,
79    dtype: _DTypeLikeFloat | None = ...,
80    order: _OrderKACF = ...,
81    casting: _CastingSafe = ...,
82    optimize: _OptimizeKind = False,
83) -> Any: ...
84@overload
85def einsum(
86    subscripts: str | _ArrayLikeInt_co,
87    /,
88    *operands: _ArrayLikeComplex_co,
89    out: None = None,
90    dtype: _DTypeLikeComplex | None = ...,
91    order: _OrderKACF = ...,
92    casting: _CastingSafe = ...,
93    optimize: _OptimizeKind = False,
94) -> Any: ...
95@overload
96def einsum(
97    subscripts: str | _ArrayLikeInt_co,
98    /,
99    *operands: Any,
100    casting: _CastingUnsafe,
101    dtype: _DTypeLikeComplex_co | None = ...,
102    out: None = None,
103    order: _OrderKACF = ...,
104    optimize: _OptimizeKind = False,
105) -> Any: ...
106@overload
107def einsum(
108    subscripts: str | _ArrayLikeInt_co,
109    /,
110    *operands: _ArrayLikeComplex_co,
111    out: _ArrayT,
112    dtype: _DTypeLikeComplex_co | None = ...,
113    order: _OrderKACF = ...,
114    casting: _CastingSafe = ...,
115    optimize: _OptimizeKind = False,
116) -> _ArrayT: ...
117@overload
118def einsum(
119    subscripts: str | _ArrayLikeInt_co,
120    /,
121    *operands: Any,
122    out: _ArrayT,
123    casting: _CastingUnsafe,
124    dtype: _DTypeLikeComplex_co | None = ...,
125    order: _OrderKACF = ...,
126    optimize: _OptimizeKind = False,
127) -> _ArrayT: ...
128
129@overload
130def einsum(
131    subscripts: str | _ArrayLikeInt_co,
132    /,
133    *operands: _ArrayLikeObject_co,
134    out: None = None,
135    dtype: _DTypeLikeObject | None = ...,
136    order: _OrderKACF = ...,
137    casting: _CastingSafe = ...,
138    optimize: _OptimizeKind = False,
139) -> Any: ...
140@overload
141def einsum(
142    subscripts: str | _ArrayLikeInt_co,
143    /,
144    *operands: Any,
145    casting: _CastingUnsafe,
146    dtype: _DTypeLikeObject | None = ...,
147    out: None = None,
148    order: _OrderKACF = ...,
149    optimize: _OptimizeKind = False,
150) -> Any: ...
151@overload
152def einsum(
153    subscripts: str | _ArrayLikeInt_co,
154    /,
155    *operands: _ArrayLikeObject_co,
156    out: _ArrayT,
157    dtype: _DTypeLikeObject | None = ...,
158    order: _OrderKACF = ...,
159    casting: _CastingSafe = ...,
160    optimize: _OptimizeKind = False,
161) -> _ArrayT: ...
162@overload
163def einsum(
164    subscripts: str | _ArrayLikeInt_co,
165    /,
166    *operands: Any,
167    out: _ArrayT,
168    casting: _CastingUnsafe,
169    dtype: _DTypeLikeObject | None = ...,
170    order: _OrderKACF = ...,
171    optimize: _OptimizeKind = False,
172) -> _ArrayT: ...
173
174# NOTE: `einsum_call` is a hidden kwarg unavailable for public use.
175# It is therefore excluded from the signatures below.
176# NOTE: In practice the list consists of a `str` (first element)
177# and a variable number of integer tuples.
178def einsum_path(
179    subscripts: str | _ArrayLikeInt_co,
180    /,
181    *operands: _ArrayLikeComplex_co | _DTypeLikeObject,
182    optimize: _OptimizeKind = "greedy",
183    einsum_call: Literal[False] = False,
184) -> tuple[list[Any], str]: ...
185