codekingpro/portable-devtools

1from typing import (
2    Any,
3    Literal as L,
4    NoReturn,
5    SupportsIndex,
6    SupportsInt,
7    TypeAlias,
8    TypeVar,
9    overload,
10)
11
12import numpy as np
13from numpy import (
14    complex128,
15    complexfloating,
16    float64,
17    floating,
18    int32,
19    int64,
20    object_,
21    poly1d,
22    signedinteger,
23    unsignedinteger,
24)
25from numpy._typing import (
26    ArrayLike,
27    NDArray,
28    _ArrayLikeBool_co,
29    _ArrayLikeComplex_co,
30    _ArrayLikeFloat_co,
31    _ArrayLikeInt_co,
32    _ArrayLikeObject_co,
33    _ArrayLikeUInt_co,
34)
35
36_T = TypeVar("_T")
37
38_2Tup: TypeAlias = tuple[_T, _T]
39_5Tup: TypeAlias = tuple[
40    _T,
41    NDArray[float64],
42    NDArray[int32],
43    NDArray[float64],
44    NDArray[float64],
45]
46
47__all__ = [
48    "poly",
49    "roots",
50    "polyint",
51    "polyder",
52    "polyadd",
53    "polysub",
54    "polymul",
55    "polydiv",
56    "polyval",
57    "poly1d",
58    "polyfit",
59]
60
61def poly(seq_of_zeros: ArrayLike) -> NDArray[floating]: ...
62
63# Returns either a float or complex array depending on the input values.
64# See `np.linalg.eigvals`.
65def roots(p: ArrayLike) -> NDArray[complexfloating] | NDArray[floating]: ...
66
67@overload
68def polyint(
69    p: poly1d,
70    m: SupportsInt | SupportsIndex = 1,
71    k: _ArrayLikeComplex_co | _ArrayLikeObject_co | None = None,
72) -> poly1d: ...
73@overload
74def polyint(
75    p: _ArrayLikeFloat_co,
76    m: SupportsInt | SupportsIndex = 1,
77    k: _ArrayLikeFloat_co | None = None,
78) -> NDArray[floating]: ...
79@overload
80def polyint(
81    p: _ArrayLikeComplex_co,
82    m: SupportsInt | SupportsIndex = 1,
83    k: _ArrayLikeComplex_co | None = None,
84) -> NDArray[complexfloating]: ...
85@overload
86def polyint(
87    p: _ArrayLikeObject_co,
88    m: SupportsInt | SupportsIndex = 1,
89    k: _ArrayLikeObject_co | None = None,
90) -> NDArray[object_]: ...
91
92@overload
93def polyder(
94    p: poly1d,
95    m: SupportsInt | SupportsIndex = 1,
96) -> poly1d: ...
97@overload
98def polyder(
99    p: _ArrayLikeFloat_co,
100    m: SupportsInt | SupportsIndex = 1,
101) -> NDArray[floating]: ...
102@overload
103def polyder(
104    p: _ArrayLikeComplex_co,
105    m: SupportsInt | SupportsIndex = 1,
106) -> NDArray[complexfloating]: ...
107@overload
108def polyder(
109    p: _ArrayLikeObject_co,
110    m: SupportsInt | SupportsIndex = 1,
111) -> NDArray[object_]: ...
112
113@overload
114def polyfit(
115    x: _ArrayLikeFloat_co,
116    y: _ArrayLikeFloat_co,
117    deg: SupportsIndex | SupportsInt,
118    rcond: float | None = None,
119    full: L[False] = False,
120    w: _ArrayLikeFloat_co | None = None,
121    cov: L[False] = False,
122) -> NDArray[float64]: ...
123@overload
124def polyfit(
125    x: _ArrayLikeComplex_co,
126    y: _ArrayLikeComplex_co,
127    deg: SupportsIndex | SupportsInt,
128    rcond: float | None = None,
129    full: L[False] = False,
130    w: _ArrayLikeFloat_co | None = None,
131    cov: L[False] = False,
132) -> NDArray[complex128]: ...
133@overload
134def polyfit(
135    x: _ArrayLikeFloat_co,
136    y: _ArrayLikeFloat_co,
137    deg: SupportsIndex | SupportsInt,
138    rcond: float | None = None,
139    full: L[False] = False,
140    w: _ArrayLikeFloat_co | None = None,
141    *,
142    cov: L[True, "unscaled"],
143) -> _2Tup[NDArray[float64]]: ...
144@overload
145def polyfit(
146    x: _ArrayLikeComplex_co,
147    y: _ArrayLikeComplex_co,
148    deg: SupportsIndex | SupportsInt,
149    rcond: float | None = None,
150    full: L[False] = False,
151    w: _ArrayLikeFloat_co | None = None,
152    *,
153    cov: L[True, "unscaled"],
154) -> _2Tup[NDArray[complex128]]: ...
155@overload
156def polyfit(
157    x: _ArrayLikeFloat_co,
158    y: _ArrayLikeFloat_co,
159    deg: SupportsIndex | SupportsInt,
160    rcond: float | None,
161    full: L[True],
162    w: _ArrayLikeFloat_co | None = None,
163    cov: bool | L["unscaled"] = False,
164) -> _5Tup[NDArray[float64]]: ...
165@overload
166def polyfit(
167    x: _ArrayLikeFloat_co,
168    y: _ArrayLikeFloat_co,
169    deg: SupportsIndex | SupportsInt,
170    rcond: float | None = None,
171    *,
172    full: L[True],
173    w: _ArrayLikeFloat_co | None = None,
174    cov: bool | L["unscaled"] = False,
175) -> _5Tup[NDArray[float64]]: ...
176@overload
177def polyfit(
178    x: _ArrayLikeComplex_co,
179    y: _ArrayLikeComplex_co,
180    deg: SupportsIndex | SupportsInt,
181    rcond: float | None,
182    full: L[True],
183    w: _ArrayLikeFloat_co | None = None,
184    cov: bool | L["unscaled"] = False,
185) -> _5Tup[NDArray[complex128]]: ...
186@overload
187def polyfit(
188    x: _ArrayLikeComplex_co,
189    y: _ArrayLikeComplex_co,
190    deg: SupportsIndex | SupportsInt,
191    rcond: float | None = None,
192    *,
193    full: L[True],
194    w: _ArrayLikeFloat_co | None = None,
195    cov: bool | L["unscaled"] = False,
196) -> _5Tup[NDArray[complex128]]: ...
197
198@overload
199def polyval(
200    p: _ArrayLikeBool_co,
201    x: _ArrayLikeBool_co,
202) -> NDArray[int64]: ...
203@overload
204def polyval(
205    p: _ArrayLikeUInt_co,
206    x: _ArrayLikeUInt_co,
207) -> NDArray[unsignedinteger]: ...
208@overload
209def polyval(
210    p: _ArrayLikeInt_co,
211    x: _ArrayLikeInt_co,
212) -> NDArray[signedinteger]: ...
213@overload
214def polyval(
215    p: _ArrayLikeFloat_co,
216    x: _ArrayLikeFloat_co,
217) -> NDArray[floating]: ...
218@overload
219def polyval(
220    p: _ArrayLikeComplex_co,
221    x: _ArrayLikeComplex_co,
222) -> NDArray[complexfloating]: ...
223@overload
224def polyval(
225    p: _ArrayLikeObject_co,
226    x: _ArrayLikeObject_co,
227) -> NDArray[object_]: ...
228
229@overload
230def polyadd(
231    a1: poly1d,
232    a2: _ArrayLikeComplex_co | _ArrayLikeObject_co,
233) -> poly1d: ...
234@overload
235def polyadd(
236    a1: _ArrayLikeComplex_co | _ArrayLikeObject_co,
237    a2: poly1d,
238) -> poly1d: ...
239@overload
240def polyadd(
241    a1: _ArrayLikeBool_co,
242    a2: _ArrayLikeBool_co,
243) -> NDArray[np.bool]: ...
244@overload
245def polyadd(
246    a1: _ArrayLikeUInt_co,
247    a2: _ArrayLikeUInt_co,
248) -> NDArray[unsignedinteger]: ...
249@overload
250def polyadd(
251    a1: _ArrayLikeInt_co,
252    a2: _ArrayLikeInt_co,
253) -> NDArray[signedinteger]: ...
254@overload
255def polyadd(
256    a1: _ArrayLikeFloat_co,
257    a2: _ArrayLikeFloat_co,
258) -> NDArray[floating]: ...
259@overload
260def polyadd(
261    a1: _ArrayLikeComplex_co,
262    a2: _ArrayLikeComplex_co,
263) -> NDArray[complexfloating]: ...
264@overload
265def polyadd(
266    a1: _ArrayLikeObject_co,
267    a2: _ArrayLikeObject_co,
268) -> NDArray[object_]: ...
269
270@overload
271def polysub(
272    a1: poly1d,
273    a2: _ArrayLikeComplex_co | _ArrayLikeObject_co,
274) -> poly1d: ...
275@overload
276def polysub(
277    a1: _ArrayLikeComplex_co | _ArrayLikeObject_co,
278    a2: poly1d,
279) -> poly1d: ...
280@overload
281def polysub(
282    a1: _ArrayLikeBool_co,
283    a2: _ArrayLikeBool_co,
284) -> NoReturn: ...
285@overload
286def polysub(
287    a1: _ArrayLikeUInt_co,
288    a2: _ArrayLikeUInt_co,
289) -> NDArray[unsignedinteger]: ...
290@overload
291def polysub(
292    a1: _ArrayLikeInt_co,
293    a2: _ArrayLikeInt_co,
294) -> NDArray[signedinteger]: ...
295@overload
296def polysub(
297    a1: _ArrayLikeFloat_co,
298    a2: _ArrayLikeFloat_co,
299) -> NDArray[floating]: ...
300@overload
301def polysub(
302    a1: _ArrayLikeComplex_co,
303    a2: _ArrayLikeComplex_co,
304) -> NDArray[complexfloating]: ...
305@overload
306def polysub(
307    a1: _ArrayLikeObject_co,
308    a2: _ArrayLikeObject_co,
309) -> NDArray[object_]: ...
310
311# NOTE: Not an alias, but they do have the same signature (that we can reuse)
312polymul = polyadd
313
314@overload
315def polydiv(
316    u: poly1d,
317    v: _ArrayLikeComplex_co | _ArrayLikeObject_co,
318) -> _2Tup[poly1d]: ...
319@overload
320def polydiv(
321    u: _ArrayLikeComplex_co | _ArrayLikeObject_co,
322    v: poly1d,
323) -> _2Tup[poly1d]: ...
324@overload
325def polydiv(
326    u: _ArrayLikeFloat_co,
327    v: _ArrayLikeFloat_co,
328) -> _2Tup[NDArray[floating]]: ...
329@overload
330def polydiv(
331    u: _ArrayLikeComplex_co,
332    v: _ArrayLikeComplex_co,
333) -> _2Tup[NDArray[complexfloating]]: ...
334@overload
335def polydiv(
336    u: _ArrayLikeObject_co,
337    v: _ArrayLikeObject_co,
338) -> _2Tup[NDArray[Any]]: ...
339