codekingpro/portable-devtools

1from collections.abc import Callable, Iterable, Sequence
2from typing import (
3    Final,
4    Literal,
5    Protocol,
6    SupportsIndex,
7    TypeAlias,
8    TypeVar,
9    overload,
10    type_check_only,
11)
12
13import numpy as np
14import numpy.typing as npt
15from numpy._typing import (
16    _ArrayLikeComplex_co,
17    _ArrayLikeFloat_co,
18    _ArrayLikeObject_co,
19    _FloatLike_co,
20    _NumberLike_co,
21)
22
23from ._polytypes import (
24    _AnyInt,
25    _Array2,
26    _ArrayLikeCoef_co,
27    _CoefArray,
28    _CoefLike_co,
29    _CoefSeries,
30    _ComplexArray,
31    _ComplexSeries,
32    _FloatArray,
33    _FloatSeries,
34    _FuncBinOp,
35    _ObjectArray,
36    _ObjectSeries,
37    _SeriesLikeCoef_co,
38    _SeriesLikeComplex_co,
39    _SeriesLikeFloat_co,
40    _SeriesLikeInt_co,
41    _SeriesLikeObject_co,
42    _Tuple2,
43)
44
45__all__ = ["as_series", "format_float", "getdomain", "mapdomain", "mapparms", "trimcoef", "trimseq"]
46
47_T = TypeVar("_T")
48_SeqT = TypeVar("_SeqT", bound=_CoefArray | Sequence[_CoefLike_co])
49
50_AnyLineF: TypeAlias = Callable[[float, float], _CoefArray]
51_AnyMulF: TypeAlias = Callable[[np.ndarray | list[int], np.ndarray], _CoefArray]
52_AnyVanderF: TypeAlias = Callable[[np.ndarray, int], _CoefArray]
53
54@type_check_only
55class _ValFunc(Protocol[_T]):
56    def __call__(self, x: np.ndarray, c: _T, /, *, tensor: bool = True) -> _T: ...
57
58###
59
60@overload
61def as_series(alist: npt.NDArray[np.integer] | _FloatArray, trim: bool = True) -> list[_FloatSeries]: ...
62@overload
63def as_series(alist: _ComplexArray, trim: bool = True) -> list[_ComplexSeries]: ...
64@overload
65def as_series(alist: _ObjectArray, trim: bool = True) -> list[_ObjectSeries]: ...
66@overload
67def as_series(alist: Iterable[_FloatArray | npt.NDArray[np.integer]], trim: bool = True) -> list[_FloatSeries]: ...
68@overload
69def as_series(alist: Iterable[_ComplexArray], trim: bool = True) -> list[_ComplexSeries]: ...
70@overload
71def as_series(alist: Iterable[_ObjectArray], trim: bool = True) -> list[_ObjectSeries]: ...
72@overload
73def as_series(alist: Iterable[_SeriesLikeFloat_co | float], trim: bool = True) -> list[_FloatSeries]: ...
74@overload
75def as_series(alist: Iterable[_SeriesLikeComplex_co | complex], trim: bool = True) -> list[_ComplexSeries]: ...
76@overload
77def as_series(alist: Iterable[_SeriesLikeCoef_co | object], trim: bool = True) -> list[_ObjectSeries]: ...
78
79#
80def trimseq(seq: _SeqT) -> _SeqT: ...
81
82#
83@overload
84def trimcoef(c: npt.NDArray[np.integer] | _FloatArray, tol: _FloatLike_co = 0) -> _FloatSeries: ...
85@overload
86def trimcoef(c: _ComplexArray, tol: _FloatLike_co = 0) -> _ComplexSeries: ...
87@overload
88def trimcoef(c: _ObjectArray, tol: _FloatLike_co = 0) -> _ObjectSeries: ...
89@overload
90def trimcoef(c: _SeriesLikeFloat_co | float, tol: _FloatLike_co = 0) -> _FloatSeries: ...
91@overload
92def trimcoef(c: _SeriesLikeComplex_co | complex, tol: _FloatLike_co = 0) -> _ComplexSeries: ...
93@overload
94def trimcoef(c: _SeriesLikeCoef_co | object, tol: _FloatLike_co = 0) -> _ObjectSeries: ...
95
96#
97@overload
98def getdomain(x: _FloatArray | npt.NDArray[np.integer]) -> _Array2[np.float64]: ...
99@overload
100def getdomain(x: _ComplexArray) -> _Array2[np.complex128]: ...
101@overload
102def getdomain(x: _ObjectArray) -> _Array2[np.object_]: ...
103@overload
104def getdomain(x: _SeriesLikeFloat_co | float) -> _Array2[np.float64]: ...
105@overload
106def getdomain(x: _SeriesLikeComplex_co | complex) -> _Array2[np.complex128]: ...
107@overload
108def getdomain(x: _SeriesLikeCoef_co | object) -> _Array2[np.object_]: ...
109
110#
111@overload
112def mapparms(old: npt.NDArray[np.floating | np.integer], new: npt.NDArray[np.floating | np.integer]) -> _Tuple2[np.floating]: ...
113@overload
114def mapparms(old: npt.NDArray[np.number], new: npt.NDArray[np.number]) -> _Tuple2[np.complexfloating]: ...
115@overload
116def mapparms(old: npt.NDArray[np.object_ | np.number], new: npt.NDArray[np.object_ | np.number]) -> _Tuple2[object]: ...
117@overload
118def mapparms(old: Sequence[float], new: Sequence[float]) -> _Tuple2[float]: ...
119@overload
120def mapparms(old: Sequence[complex], new: Sequence[complex]) -> _Tuple2[complex]: ...
121@overload
122def mapparms(old: _SeriesLikeFloat_co, new: _SeriesLikeFloat_co) -> _Tuple2[np.floating]: ...
123@overload
124def mapparms(old: _SeriesLikeComplex_co, new: _SeriesLikeComplex_co) -> _Tuple2[np.complexfloating]: ...
125@overload
126def mapparms(old: _SeriesLikeCoef_co, new: _SeriesLikeCoef_co) -> _Tuple2[object]: ...
127
128#
129@overload
130def mapdomain(x: _FloatLike_co, old: _SeriesLikeFloat_co, new: _SeriesLikeFloat_co) -> np.floating: ...
131@overload
132def mapdomain(x: _NumberLike_co, old: _SeriesLikeComplex_co, new: _SeriesLikeComplex_co) -> np.complexfloating: ...
133@overload
134def mapdomain(
135    x: npt.NDArray[np.floating | np.integer],
136    old: npt.NDArray[np.floating | np.integer],
137    new: npt.NDArray[np.floating | np.integer],
138) -> _FloatSeries: ...
139@overload
140def mapdomain(x: npt.NDArray[np.number], old: npt.NDArray[np.number], new: npt.NDArray[np.number]) -> _ComplexSeries: ...
141@overload
142def mapdomain(
143    x: npt.NDArray[np.object_ | np.number],
144    old: npt.NDArray[np.object_ | np.number],
145    new: npt.NDArray[np.object_ | np.number],
146) -> _ObjectSeries: ...
147@overload
148def mapdomain(x: _SeriesLikeFloat_co, old: _SeriesLikeFloat_co, new: _SeriesLikeFloat_co) -> _FloatSeries: ...
149@overload
150def mapdomain(x: _SeriesLikeComplex_co, old: _SeriesLikeComplex_co, new: _SeriesLikeComplex_co) -> _ComplexSeries: ...
151@overload
152def mapdomain(x: _SeriesLikeCoef_co, old: _SeriesLikeCoef_co, new: _SeriesLikeCoef_co) -> _ObjectSeries: ...
153@overload
154def mapdomain(x: _CoefLike_co, old: _SeriesLikeCoef_co, new: _SeriesLikeCoef_co) -> object: ...
155
156#
157def _nth_slice(i: SupportsIndex, ndim: SupportsIndex) -> tuple[slice | None, ...]: ...
158
159# keep in sync with `vander_nd_flat`
160@overload
161def _vander_nd(
162    vander_fs: Sequence[_AnyVanderF],
163    points: Sequence[_ArrayLikeFloat_co],
164    degrees: Sequence[SupportsIndex],
165) -> _FloatArray: ...
166@overload
167def _vander_nd(
168    vander_fs: Sequence[_AnyVanderF],
169    points: Sequence[_ArrayLikeComplex_co],
170    degrees: Sequence[SupportsIndex],
171) -> _ComplexArray: ...
172@overload
173def _vander_nd(
174    vander_fs: Sequence[_AnyVanderF],
175    points: Sequence[_ArrayLikeObject_co | _ArrayLikeComplex_co],
176    degrees: Sequence[SupportsIndex],
177) -> _ObjectArray: ...
178@overload
179def _vander_nd(
180    vander_fs: Sequence[_AnyVanderF],
181    points: Sequence[npt.ArrayLike],
182    degrees: Sequence[SupportsIndex],
183) -> _CoefArray: ...
184
185# keep in sync with `vander_nd`
186@overload
187def _vander_nd_flat(
188    vander_fs: Sequence[_AnyVanderF],
189    points: Sequence[_ArrayLikeFloat_co],
190    degrees: Sequence[SupportsIndex],
191) -> _FloatArray: ...
192@overload
193def _vander_nd_flat(
194    vander_fs: Sequence[_AnyVanderF],
195    points: Sequence[_ArrayLikeComplex_co],
196    degrees: Sequence[SupportsIndex],
197) -> _ComplexArray: ...
198@overload
199def _vander_nd_flat(
200    vander_fs: Sequence[_AnyVanderF],
201    points: Sequence[_ArrayLikeObject_co | _ArrayLikeComplex_co],
202    degrees: Sequence[SupportsIndex],
203) -> _ObjectArray: ...
204@overload
205def _vander_nd_flat(
206    vander_fs: Sequence[_AnyVanderF],
207    points: Sequence[npt.ArrayLike],
208    degrees: Sequence[SupportsIndex],
209) -> _CoefArray: ...
210
211# keep in sync with `._polytypes._FuncFromRoots`
212@overload
213def _fromroots(line_f: _AnyLineF, mul_f: _AnyMulF, roots: _SeriesLikeFloat_co) -> _FloatSeries: ...
214@overload
215def _fromroots(line_f: _AnyLineF, mul_f: _AnyMulF, roots: _SeriesLikeComplex_co) -> _ComplexSeries: ...
216@overload
217def _fromroots(line_f: _AnyLineF, mul_f: _AnyMulF, roots: _SeriesLikeObject_co) -> _ObjectSeries: ...
218@overload
219def _fromroots(line_f: _AnyLineF, mul_f: _AnyMulF, roots: _SeriesLikeCoef_co) -> _CoefSeries: ...
220
221# keep in sync with `_gridnd`
222def _valnd(val_f: _ValFunc[_T], c: _T, *args: npt.ArrayLike) -> _T: ...
223
224# keep in sync with `_valnd`
225def _gridnd(val_f: _ValFunc[_T], c: _T, *args: npt.ArrayLike) -> _T: ...
226
227# keep in sync with `_polytypes._FuncBinOp`
228@overload
229def _div(mul_f: _AnyMulF, c1: _SeriesLikeFloat_co, c2: _SeriesLikeFloat_co) -> _Tuple2[_FloatSeries]: ...
230@overload
231def _div(mul_f: _AnyMulF, c1: _SeriesLikeComplex_co, c2: _SeriesLikeComplex_co) -> _Tuple2[_ComplexSeries]: ...
232@overload
233def _div(mul_f: _AnyMulF, c1: _SeriesLikeObject_co, c2: _SeriesLikeObject_co) -> _Tuple2[_ObjectSeries]: ...
234@overload
235def _div(mul_f: _AnyMulF, c1: _SeriesLikeCoef_co, c2: _SeriesLikeCoef_co) -> _Tuple2[_CoefSeries]: ...
236
237_add: Final[_FuncBinOp] = ...
238_sub: Final[_FuncBinOp] = ...
239
240# keep in sync with `_polytypes._FuncPow`
241@overload
242def _pow(mul_f: _AnyMulF, c: _SeriesLikeFloat_co, pow: _AnyInt, maxpower: _AnyInt | None) -> _FloatSeries: ...
243@overload
244def _pow(mul_f: _AnyMulF, c: _SeriesLikeComplex_co, pow: _AnyInt, maxpower: _AnyInt | None) -> _ComplexSeries: ...
245@overload
246def _pow(mul_f: _AnyMulF, c: _SeriesLikeObject_co, pow: _AnyInt, maxpower: _AnyInt | None) -> _ObjectSeries: ...
247@overload
248def _pow(mul_f: _AnyMulF, c: _SeriesLikeCoef_co, pow: _AnyInt, maxpower: _AnyInt | None) -> _CoefSeries: ...
249
250# keep in sync with `_polytypes._FuncFit`
251@overload
252def _fit(
253    vander_f: _AnyVanderF,
254    x: _SeriesLikeFloat_co,
255    y: _ArrayLikeFloat_co,
256    deg: _SeriesLikeInt_co,
257    rcond: _FloatLike_co | None = None,
258    full: Literal[False] = False,
259    w: _SeriesLikeFloat_co | None = None,
260) -> _FloatArray: ...
261@overload
262def _fit(
263    vander_f: _AnyVanderF,
264    x: _SeriesLikeComplex_co,
265    y: _ArrayLikeComplex_co,
266    deg: _SeriesLikeInt_co,
267    rcond: _FloatLike_co | None = None,
268    full: Literal[False] = False,
269    w: _SeriesLikeComplex_co | None = None,
270) -> _ComplexArray: ...
271@overload
272def _fit(
273    vander_f: _AnyVanderF,
274    x: _SeriesLikeCoef_co,
275    y: _ArrayLikeCoef_co,
276    deg: _SeriesLikeInt_co,
277    rcond: _FloatLike_co | None = None,
278    full: Literal[False] = False,
279    w: _SeriesLikeCoef_co | None = None,
280) -> _CoefArray: ...
281@overload
282def _fit(
283    vander_f: _AnyVanderF,
284    x: _SeriesLikeCoef_co,
285    y: _SeriesLikeCoef_co,
286    deg: _SeriesLikeInt_co,
287    rcond: _FloatLike_co | None,
288    full: Literal[True],
289    w: _SeriesLikeCoef_co | None = None,
290) -> tuple[_CoefSeries, Sequence[np.inexact | np.int32]]: ...
291@overload
292def _fit(
293    vander_f: _AnyVanderF,
294    x: _SeriesLikeCoef_co,
295    y: _SeriesLikeCoef_co,
296    deg: _SeriesLikeInt_co,
297    rcond: _FloatLike_co | None = None,
298    *,
299    full: Literal[True],
300    w: _SeriesLikeCoef_co | None = None,
301) -> tuple[_CoefSeries, Sequence[np.inexact | np.int32]]: ...
302
303#
304def _as_int(x: SupportsIndex, desc: str) -> int: ...
305
306#
307def format_float(x: _FloatLike_co, parens: bool = False) -> str: ...
308