codekingpro/portable-devtools

1from typing import (
2    Literal,
3    Never,
4    TypeAlias,
5    TypedDict,
6    TypeVar,
7    final,
8    overload,
9    type_check_only,
10)
11
12import numpy as np
13
14_Device: TypeAlias = Literal["cpu"]
15_DeviceLike: TypeAlias = _Device | None
16
17_Capabilities = TypedDict(
18    "_Capabilities",
19    {
20        "boolean indexing": Literal[True],
21        "data-dependent shapes": Literal[True],
22    },
23)
24
25_DefaultDTypes = TypedDict(
26    "_DefaultDTypes",
27    {
28        "real floating": np.dtype[np.float64],
29        "complex floating": np.dtype[np.complex128],
30        "integral": np.dtype[np.intp],
31        "indexing": np.dtype[np.intp],
32    },
33)
34
35_KindBool: TypeAlias = Literal["bool"]
36_KindInt: TypeAlias = Literal["signed integer"]
37_KindUInt: TypeAlias = Literal["unsigned integer"]
38_KindInteger: TypeAlias = Literal["integral"]
39_KindFloat: TypeAlias = Literal["real floating"]
40_KindComplex: TypeAlias = Literal["complex floating"]
41_KindNumber: TypeAlias = Literal["numeric"]
42_Kind: TypeAlias = (
43    _KindBool
44    | _KindInt
45    | _KindUInt
46    | _KindInteger
47    | _KindFloat
48    | _KindComplex
49    | _KindNumber
50)
51
52_T1 = TypeVar("_T1")
53_T2 = TypeVar("_T2")
54_T3 = TypeVar("_T3")
55_Permute1: TypeAlias = _T1 | tuple[_T1]
56_Permute2: TypeAlias = tuple[_T1, _T2] | tuple[_T2, _T1]
57_Permute3: TypeAlias = (
58    tuple[_T1, _T2, _T3] | tuple[_T1, _T3, _T2]
59    | tuple[_T2, _T1, _T3] | tuple[_T2, _T3, _T1]
60    | tuple[_T3, _T1, _T2] | tuple[_T3, _T2, _T1]
61)
62
63@type_check_only
64class _DTypesBool(TypedDict):
65    bool: np.dtype[np.bool]
66
67@type_check_only
68class _DTypesInt(TypedDict):
69    int8: np.dtype[np.int8]
70    int16: np.dtype[np.int16]
71    int32: np.dtype[np.int32]
72    int64: np.dtype[np.int64]
73
74@type_check_only
75class _DTypesUInt(TypedDict):
76    uint8: np.dtype[np.uint8]
77    uint16: np.dtype[np.uint16]
78    uint32: np.dtype[np.uint32]
79    uint64: np.dtype[np.uint64]
80
81@type_check_only
82class _DTypesInteger(_DTypesInt, _DTypesUInt): ...
83
84@type_check_only
85class _DTypesFloat(TypedDict):
86    float32: np.dtype[np.float32]
87    float64: np.dtype[np.float64]
88
89@type_check_only
90class _DTypesComplex(TypedDict):
91    complex64: np.dtype[np.complex64]
92    complex128: np.dtype[np.complex128]
93
94@type_check_only
95class _DTypesNumber(_DTypesInteger, _DTypesFloat, _DTypesComplex): ...
96
97@type_check_only
98class _DTypes(_DTypesBool, _DTypesNumber): ...
99
100@type_check_only
101class _DTypesUnion(TypedDict, total=False):
102    bool: np.dtype[np.bool]
103    int8: np.dtype[np.int8]
104    int16: np.dtype[np.int16]
105    int32: np.dtype[np.int32]
106    int64: np.dtype[np.int64]
107    uint8: np.dtype[np.uint8]
108    uint16: np.dtype[np.uint16]
109    uint32: np.dtype[np.uint32]
110    uint64: np.dtype[np.uint64]
111    float32: np.dtype[np.float32]
112    float64: np.dtype[np.float64]
113    complex64: np.dtype[np.complex64]
114    complex128: np.dtype[np.complex128]
115
116_EmptyDict: TypeAlias = dict[Never, Never]
117
118@final
119class __array_namespace_info__:
120    __module__: Literal["numpy"] = "numpy"
121
122    def capabilities(self) -> _Capabilities: ...
123    def default_device(self) -> _Device: ...
124    def default_dtypes(
125        self,
126        *,
127        device: _DeviceLike = None,
128    ) -> _DefaultDTypes: ...
129    def devices(self) -> list[_Device]: ...
130
131    @overload
132    def dtypes(
133        self,
134        *,
135        device: _DeviceLike = None,
136        kind: None = None,
137    ) -> _DTypes: ...
138    @overload
139    def dtypes(
140        self,
141        *,
142        device: _DeviceLike = None,
143        kind: _Permute1[_KindBool],
144    ) -> _DTypesBool: ...
145    @overload
146    def dtypes(
147        self,
148        *,
149        device: _DeviceLike = None,
150        kind: _Permute1[_KindInt],
151    ) -> _DTypesInt: ...
152    @overload
153    def dtypes(
154        self,
155        *,
156        device: _DeviceLike = None,
157        kind: _Permute1[_KindUInt],
158    ) -> _DTypesUInt: ...
159    @overload
160    def dtypes(
161        self,
162        *,
163        device: _DeviceLike = None,
164        kind: _Permute1[_KindFloat],
165    ) -> _DTypesFloat: ...
166    @overload
167    def dtypes(
168        self,
169        *,
170        device: _DeviceLike = None,
171        kind: _Permute1[_KindComplex],
172    ) -> _DTypesComplex: ...
173    @overload
174    def dtypes(
175        self,
176        *,
177        device: _DeviceLike = None,
178        kind: (
179            _Permute1[_KindInteger]
180            | _Permute2[_KindInt, _KindUInt]
181        ),
182    ) -> _DTypesInteger: ...
183    @overload
184    def dtypes(
185        self,
186        *,
187        device: _DeviceLike = None,
188        kind: (
189            _Permute1[_KindNumber]
190            | _Permute3[_KindInteger, _KindFloat, _KindComplex]
191        ),
192    ) -> _DTypesNumber: ...
193    @overload
194    def dtypes(
195        self,
196        *,
197        device: _DeviceLike = None,
198        kind: tuple[()],
199    ) -> _EmptyDict: ...
200    @overload
201    def dtypes(
202        self,
203        *,
204        device: _DeviceLike = None,
205        kind: tuple[_Kind, ...],
206    ) -> _DTypesUnion: ...
207