codekingpro/portable-devtools

1from _typeshed import Incomplete
2from collections.abc import Callable, Sequence
3from typing import (
4    Any,
5    Literal as L,
6    Never,
7    Protocol,
8    TypeAlias,
9    TypeVar,
10    overload,
11    type_check_only,
12)
13
14import numpy as np
15from numpy import _OrderCF
16from numpy._typing import (
17    ArrayLike,
18    DTypeLike,
19    NDArray,
20    _ArrayLike,
21    _DTypeLike,
22    _NumberLike_co,
23    _ScalarLike_co,
24    _SupportsArray,
25    _SupportsArrayFunc,
26)
27
28__all__ = [
29    "diag",
30    "diagflat",
31    "eye",
32    "fliplr",
33    "flipud",
34    "tri",
35    "triu",
36    "tril",
37    "vander",
38    "histogram2d",
39    "mask_indices",
40    "tril_indices",
41    "tril_indices_from",
42    "triu_indices",
43    "triu_indices_from",
44]
45
46###
47
48_T = TypeVar("_T")
49_ArrayT = TypeVar("_ArrayT", bound=np.ndarray)
50_ScalarT = TypeVar("_ScalarT", bound=np.generic)
51_ComplexT = TypeVar("_ComplexT", bound=np.complexfloating)
52_InexactT = TypeVar("_InexactT", bound=np.inexact)
53_NumberT = TypeVar("_NumberT", bound=np.number)
54_NumberObjectT = TypeVar("_NumberObjectT", bound=np.number | np.object_)
55_NumberCoT = TypeVar("_NumberCoT", bound=_Number_co)
56
57_Int_co: TypeAlias = np.integer | np.bool
58_Float_co: TypeAlias = np.floating | _Int_co
59_Number_co: TypeAlias = np.number | np.bool
60
61_Array1D: TypeAlias = np.ndarray[tuple[int], np.dtype[_ScalarT]]
62_Array2D: TypeAlias = np.ndarray[tuple[int, int], np.dtype[_ScalarT]]
63# Workaround for mypy's and pyright's lack of compliance with the typing spec for
64# overloads for gradual types. This works because only `Any` and `Never` are assignable
65# to `Never`.
66_ArrayNoD: TypeAlias = np.ndarray[tuple[Never] | tuple[Never, Never], np.dtype[_ScalarT]]
67
68_ArrayLike1D: TypeAlias = _SupportsArray[np.dtype[_ScalarT]] | Sequence[_ScalarT]
69_ArrayLike1DInt_co: TypeAlias = _SupportsArray[np.dtype[_Int_co]] | Sequence[int | _Int_co]
70_ArrayLike1DFloat_co: TypeAlias = _SupportsArray[np.dtype[_Float_co]] | Sequence[float | _Float_co]
71_ArrayLike2DFloat_co: TypeAlias = _SupportsArray[np.dtype[_Float_co]] | Sequence[_ArrayLike1DFloat_co]
72_ArrayLike1DNumber_co: TypeAlias = _SupportsArray[np.dtype[_Number_co]] | Sequence[complex | _Number_co]
73
74# The returned arrays dtype must be compatible with `np.equal`
75_MaskFunc: TypeAlias = Callable[[NDArray[np.int_], _T], NDArray[_Number_co | np.timedelta64 | np.datetime64 | np.object_]]
76
77_Indices2D: TypeAlias = tuple[_Array1D[np.intp], _Array1D[np.intp]]
78_Histogram2D: TypeAlias = tuple[_Array2D[np.float64], _Array1D[_ScalarT], _Array1D[_ScalarT]]
79
80@type_check_only
81class _HasShapeAndNDim(Protocol):
82    @property  # TODO: require 2d shape once shape-typing has matured
83    def shape(self) -> tuple[int, ...]: ...
84    @property
85    def ndim(self) -> int: ...
86
87###
88
89# keep in sync with `flipud`
90@overload
91def fliplr(m: _ArrayT) -> _ArrayT: ...
92@overload
93def fliplr(m: _ArrayLike[_ScalarT]) -> NDArray[_ScalarT]: ...
94@overload
95def fliplr(m: ArrayLike) -> NDArray[Any]: ...
96
97# keep in sync with `fliplr`
98@overload
99def flipud(m: _ArrayT) -> _ArrayT: ...
100@overload
101def flipud(m: _ArrayLike[_ScalarT]) -> NDArray[_ScalarT]: ...
102@overload
103def flipud(m: ArrayLike) -> NDArray[Any]: ...
104
105#
106@overload
107def eye(
108    N: int,
109    M: int | None = None,
110    k: int = 0,
111    dtype: None = ...,  # = float  # stubdefaulter: ignore[missing-default]
112    order: _OrderCF = "C",
113    *,
114    device: L["cpu"] | None = None,
115    like: _SupportsArrayFunc | None = None,
116) -> _Array2D[np.float64]: ...
117@overload
118def eye(
119    N: int,
120    M: int | None,
121    k: int,
122    dtype: _DTypeLike[_ScalarT],
123    order: _OrderCF = "C",
124    *,
125    device: L["cpu"] | None = None,
126    like: _SupportsArrayFunc | None = None,
127) -> _Array2D[_ScalarT]: ...
128@overload
129def eye(
130    N: int,
131    M: int | None = None,
132    k: int = 0,
133    *,
134    dtype: _DTypeLike[_ScalarT],
135    order: _OrderCF = "C",
136    device: L["cpu"] | None = None,
137    like: _SupportsArrayFunc | None = None,
138) -> _Array2D[_ScalarT]: ...
139@overload
140def eye(
141    N: int,
142    M: int | None = None,
143    k: int = 0,
144    dtype: DTypeLike | None = ...,  # = float
145    order: _OrderCF = "C",
146    *,
147    device: L["cpu"] | None = None,
148    like: _SupportsArrayFunc | None = None,
149) -> _Array2D[Incomplete]: ...
150
151#
152@overload
153def diag(v: _ArrayNoD[_ScalarT] | Sequence[Sequence[_ScalarT]], k: int = 0) -> NDArray[_ScalarT]: ...
154@overload
155def diag(v: _Array2D[_ScalarT] | Sequence[Sequence[_ScalarT]], k: int = 0) -> _Array1D[_ScalarT]: ...
156@overload
157def diag(v: _Array1D[_ScalarT] | Sequence[_ScalarT], k: int = 0) -> _Array2D[_ScalarT]: ...
158@overload
159def diag(v: Sequence[Sequence[_ScalarLike_co]], k: int = 0) -> _Array1D[Incomplete]: ...
160@overload
161def diag(v: Sequence[_ScalarLike_co], k: int = 0) -> _Array2D[Incomplete]: ...
162@overload
163def diag(v: _ArrayLike[_ScalarT], k: int = 0) -> NDArray[_ScalarT]: ...
164@overload
165def diag(v: ArrayLike, k: int = 0) -> NDArray[Incomplete]: ...
166
167# keep in sync with `numpy.ma.extras.diagflat`
168@overload
169def diagflat(v: _ArrayLike[_ScalarT], k: int = 0) -> _Array2D[_ScalarT]: ...
170@overload
171def diagflat(v: ArrayLike, k: int = 0) -> _Array2D[Incomplete]: ...
172
173#
174@overload
175def tri(
176    N: int,
177    M: int | None = None,
178    k: int = 0,
179    dtype: None = ...,  # = float  # stubdefaulter: ignore[missing-default]
180    *,
181    like: _SupportsArrayFunc | None = None
182) -> _Array2D[np.float64]: ...
183@overload
184def tri(
185    N: int,
186    M: int | None,
187    k: int,
188    dtype: _DTypeLike[_ScalarT],
189    *,
190    like: _SupportsArrayFunc | None = None
191) -> _Array2D[_ScalarT]: ...
192@overload
193def tri(
194    N: int,
195    M: int | None = None,
196    k: int = 0,
197    *,
198    dtype: _DTypeLike[_ScalarT],
199    like: _SupportsArrayFunc | None = None
200) -> _Array2D[_ScalarT]: ...
201@overload
202def tri(
203    N: int,
204    M: int | None = None,
205    k: int = 0,
206    dtype: DTypeLike | None = ...,  # = float
207    *,
208    like: _SupportsArrayFunc | None = None
209) -> _Array2D[Any]: ...
210
211# keep in sync with `triu`
212@overload
213def tril(m: _ArrayT, k: int = 0) -> _ArrayT: ...
214@overload
215def tril(m: _ArrayLike[_ScalarT], k: int = 0) -> NDArray[_ScalarT]: ...
216@overload
217def tril(m: ArrayLike, k: int = 0) -> NDArray[Any]: ...
218
219# keep in sync with `tril`
220@overload
221def triu(m: _ArrayT, k: int = 0) -> _ArrayT: ...
222@overload
223def triu(m: _ArrayLike[_ScalarT], k: int = 0) -> NDArray[_ScalarT]: ...
224@overload
225def triu(m: ArrayLike, k: int = 0) -> NDArray[Any]: ...
226
227# we use `list` (invariant) instead of `Sequence` (covariant) to avoid overlap
228@overload
229def vander(x: _ArrayLike1D[_NumberObjectT], N: int | None = None, increasing: bool = False) -> _Array2D[_NumberObjectT]: ...
230@overload
231def vander(x: _ArrayLike1D[np.bool] | list[int], N: int | None = None, increasing: bool = False) -> _Array2D[np.int_]: ...
232@overload
233def vander(x: list[float], N: int | None = None, increasing: bool = False) -> _Array2D[np.float64]: ...
234@overload
235def vander(x: list[complex], N: int | None = None, increasing: bool = False) -> _Array2D[np.complex128]: ...
236@overload  # fallback
237def vander(x: Sequence[_NumberLike_co], N: int | None = None, increasing: bool = False) -> _Array2D[Any]: ...
238
239#
240@overload
241def histogram2d(
242    x: _ArrayLike1D[_ComplexT],
243    y: _ArrayLike1D[_ComplexT | _Float_co],
244    bins: int | Sequence[int] = 10,
245    range: _ArrayLike2DFloat_co | None = None,
246    density: bool | None = None,
247    weights: _ArrayLike1DFloat_co | None = None,
248) -> _Histogram2D[_ComplexT]: ...
249@overload
250def histogram2d(
251    x: _ArrayLike1D[_ComplexT | _Float_co],
252    y: _ArrayLike1D[_ComplexT],
253    bins: int | Sequence[int] = 10,
254    range: _ArrayLike2DFloat_co | None = None,
255    density: bool | None = None,
256    weights: _ArrayLike1DFloat_co | None = None,
257) -> _Histogram2D[_ComplexT]: ...
258@overload
259def histogram2d(
260    x: _ArrayLike1D[_InexactT],
261    y: _ArrayLike1D[_InexactT | _Int_co],
262    bins: int | Sequence[int] = 10,
263    range: _ArrayLike2DFloat_co | None = None,
264    density: bool | None = None,
265    weights: _ArrayLike1DFloat_co | None = None,
266) -> _Histogram2D[_InexactT]: ...
267@overload
268def histogram2d(
269    x: _ArrayLike1D[_InexactT | _Int_co],
270    y: _ArrayLike1D[_InexactT],
271    bins: int | Sequence[int] = 10,
272    range: _ArrayLike2DFloat_co | None = None,
273    density: bool | None = None,
274    weights: _ArrayLike1DFloat_co | None = None,
275) -> _Histogram2D[_InexactT]: ...
276@overload
277def histogram2d(
278    x: _ArrayLike1DInt_co | Sequence[float],
279    y: _ArrayLike1DInt_co | Sequence[float],
280    bins: int | Sequence[int] = 10,
281    range: _ArrayLike2DFloat_co | None = None,
282    density: bool | None = None,
283    weights: _ArrayLike1DFloat_co | None = None,
284) -> _Histogram2D[np.float64]: ...
285@overload
286def histogram2d(
287    x: Sequence[complex],
288    y: Sequence[complex],
289    bins: int | Sequence[int] = 10,
290    range: _ArrayLike2DFloat_co | None = None,
291    density: bool | None = None,
292    weights: _ArrayLike1DFloat_co | None = None,
293) -> _Histogram2D[np.complex128 | Any]: ...
294@overload
295def histogram2d(
296    x: _ArrayLike1DNumber_co,
297    y: _ArrayLike1DNumber_co,
298    bins: _ArrayLike1D[_NumberCoT] | Sequence[_ArrayLike1D[_NumberCoT]],
299    range: _ArrayLike2DFloat_co | None = None,
300    density: bool | None = None,
301    weights: _ArrayLike1DFloat_co | None = None,
302) -> _Histogram2D[_NumberCoT]: ...
303@overload
304def histogram2d(
305    x: _ArrayLike1D[_InexactT],
306    y: _ArrayLike1D[_InexactT],
307    bins: Sequence[_ArrayLike1D[_NumberCoT] | int],
308    range: _ArrayLike2DFloat_co | None = None,
309    density: bool | None = None,
310    weights: _ArrayLike1DFloat_co | None = None,
311) -> _Histogram2D[_InexactT | _NumberCoT]: ...
312@overload
313def histogram2d(
314    x: _ArrayLike1D[_InexactT],
315    y: _ArrayLike1D[_InexactT],
316    bins: Sequence[_ArrayLike1DNumber_co | int],
317    range: _ArrayLike2DFloat_co | None = None,
318    density: bool | None = None,
319    weights: _ArrayLike1DFloat_co | None = None,
320) -> _Histogram2D[_InexactT | Any]: ...
321@overload
322def histogram2d(
323    x: _ArrayLike1DInt_co | Sequence[float],
324    y: _ArrayLike1DInt_co | Sequence[float],
325    bins: Sequence[_ArrayLike1D[_NumberCoT] | int],
326    range: _ArrayLike2DFloat_co | None = None,
327    density: bool | None = None,
328    weights: _ArrayLike1DFloat_co | None = None,
329) -> _Histogram2D[np.float64 | _NumberCoT]: ...
330@overload
331def histogram2d(
332    x: _ArrayLike1DInt_co | Sequence[float],
333    y: _ArrayLike1DInt_co | Sequence[float],
334    bins: Sequence[_ArrayLike1DNumber_co | int],
335    range: _ArrayLike2DFloat_co | None = None,
336    density: bool | None = None,
337    weights: _ArrayLike1DFloat_co | None = None,
338) -> _Histogram2D[np.float64 | Any]: ...
339@overload
340def histogram2d(
341    x: Sequence[complex],
342    y: Sequence[complex],
343    bins: Sequence[_ArrayLike1D[_NumberCoT] | int],
344    range: _ArrayLike2DFloat_co | None = None,
345    density: bool | None = None,
346    weights: _ArrayLike1DFloat_co | None = None,
347) -> _Histogram2D[np.complex128 | _NumberCoT]: ...
348@overload
349def histogram2d(
350    x: Sequence[complex],
351    y: Sequence[complex],
352    bins: Sequence[_ArrayLike1DNumber_co | int],
353    range: _ArrayLike2DFloat_co | None = None,
354    density: bool | None = None,
355    weights: _ArrayLike1DFloat_co | None = None,
356) -> _Histogram2D[np.complex128 | Any]: ...
357@overload
358def histogram2d(
359    x: _ArrayLike1DNumber_co,
360    y: _ArrayLike1DNumber_co,
361    bins: Sequence[Sequence[int]],
362    range: _ArrayLike2DFloat_co | None = None,
363    density: bool | None = None,
364    weights: _ArrayLike1DFloat_co | None = None,
365) -> _Histogram2D[np.int_]: ...
366@overload
367def histogram2d(
368    x: _ArrayLike1DNumber_co,
369    y: _ArrayLike1DNumber_co,
370    bins: Sequence[Sequence[float]],
371    range: _ArrayLike2DFloat_co | None = None,
372    density: bool | None = None,
373    weights: _ArrayLike1DFloat_co | None = None,
374) -> _Histogram2D[np.float64 | Any]: ...
375@overload
376def histogram2d(
377    x: _ArrayLike1DNumber_co,
378    y: _ArrayLike1DNumber_co,
379    bins: Sequence[Sequence[complex]],
380    range: _ArrayLike2DFloat_co | None = None,
381    density: bool | None = None,
382    weights: _ArrayLike1DFloat_co | None = None,
383) -> _Histogram2D[np.complex128 | Any]: ...
384@overload
385def histogram2d(
386    x: _ArrayLike1DNumber_co,
387    y: _ArrayLike1DNumber_co,
388    bins: Sequence[_ArrayLike1DNumber_co | int] | int,
389    range: _ArrayLike2DFloat_co | None = None,
390    density: bool | None = None,
391    weights: _ArrayLike1DFloat_co | None = None,
392) -> _Histogram2D[Any]: ...
393
394# NOTE: we're assuming/demanding here the `mask_func` returns
395# an ndarray of shape `(n, n)`; otherwise there is the possibility
396# of the output tuple having more or less than 2 elements
397@overload
398def mask_indices(n: int, mask_func: _MaskFunc[int], k: int = 0) -> _Indices2D: ...
399@overload
400def mask_indices(n: int, mask_func: _MaskFunc[_T], k: _T) -> _Indices2D: ...
401
402#
403def tril_indices(n: int, k: int = 0, m: int | None = None) -> _Indices2D: ...
404def triu_indices(n: int, k: int = 0, m: int | None = None) -> _Indices2D: ...
405
406# these will accept anything with `shape: tuple[int, int]` and `ndim: int` attributes
407def tril_indices_from(arr: _HasShapeAndNDim, k: int = 0) -> _Indices2D: ...
408def triu_indices_from(arr: _HasShapeAndNDim, k: int = 0) -> _Indices2D: ...
409