1"""Private counterpart of ``numpy.typing``."""
2
3from __future__ import annotations
4
5from numpy import ufunc
6from numpy.core.overrides import set_module
7from typing import TYPE_CHECKING, final
8
9
10@final # Disallow the creation of arbitrary `NBitBase` subclasses
11@set_module("numpy.typing")
12class NBitBase:
13 """
14 A type representing `numpy.number` precision during static type checking.
15
16 Used exclusively for the purpose static type checking, `NBitBase`
17 represents the base of a hierarchical set of subclasses.
18 Each subsequent subclass is herein used for representing a lower level
19 of precision, *e.g.* ``64Bit > 32Bit > 16Bit``.
20
21 .. versionadded:: 1.20
22
23 Examples
24 --------
25 Below is a typical usage example: `NBitBase` is herein used for annotating
26 a function that takes a float and integer of arbitrary precision
27 as arguments and returns a new float of whichever precision is largest
28 (*e.g.* ``np.float16 + np.int64 -> np.float64``).
29
30 .. code-block:: python
31
32 >>> from __future__ import annotations
33 >>> from typing import TypeVar, TYPE_CHECKING
34 >>> import numpy as np
35 >>> import numpy.typing as npt
36
37 >>> T1 = TypeVar("T1", bound=npt.NBitBase)
38 >>> T2 = TypeVar("T2", bound=npt.NBitBase)
39
40 >>> def add(a: np.floating[T1], b: np.integer[T2]) -> np.floating[T1 | T2]:
41 ... return a + b
42
43 >>> a = np.float16()
44 >>> b = np.int64()
45 >>> out = add(a, b)
46
47 >>> if TYPE_CHECKING:
48 ... reveal_locals()
49 ... # note: Revealed local types are:
50 ... # note: a: numpy.floating[numpy.typing._16Bit*]
51 ... # note: b: numpy.signedinteger[numpy.typing._64Bit*]
52 ... # note: out: numpy.floating[numpy.typing._64Bit*]
53
54 """
55
56 def __init_subclass__(cls) -> None:
57 allowed_names = {
58 "NBitBase", "_256Bit", "_128Bit", "_96Bit", "_80Bit",
59 "_64Bit", "_32Bit", "_16Bit", "_8Bit",
60 }
61 if cls.__name__ not in allowed_names:
62 raise TypeError('cannot inherit from final class "NBitBase"')
63 super().__init_subclass__()
64
65
66# Silence errors about subclassing a `@final`-decorated class
67class _256Bit(NBitBase): # type: ignore[misc]
68 pass
69
70class _128Bit(_256Bit): # type: ignore[misc]
71 pass
72
73class _96Bit(_128Bit): # type: ignore[misc]
74 pass
75
76class _80Bit(_96Bit): # type: ignore[misc]
77 pass
78
79class _64Bit(_80Bit): # type: ignore[misc]
80 pass
81
82class _32Bit(_64Bit): # type: ignore[misc]
83 pass
84
85class _16Bit(_32Bit): # type: ignore[misc]
86 pass
87
88class _8Bit(_16Bit): # type: ignore[misc]
89 pass
90
91
92from ._nested_sequence import (
93 _NestedSequence as _NestedSequence,
94)
95from ._nbit import (
96 _NBitByte as _NBitByte,
97 _NBitShort as _NBitShort,
98 _NBitIntC as _NBitIntC,
99 _NBitIntP as _NBitIntP,
100 _NBitInt as _NBitInt,
101 _NBitLongLong as _NBitLongLong,
102 _NBitHalf as _NBitHalf,
103 _NBitSingle as _NBitSingle,
104 _NBitDouble as _NBitDouble,
105 _NBitLongDouble as _NBitLongDouble,
106)
107from ._char_codes import (
108 _BoolCodes as _BoolCodes,
109 _UInt8Codes as _UInt8Codes,
110 _UInt16Codes as _UInt16Codes,
111 _UInt32Codes as _UInt32Codes,
112 _UInt64Codes as _UInt64Codes,
113 _Int8Codes as _Int8Codes,
114 _Int16Codes as _Int16Codes,
115 _Int32Codes as _Int32Codes,
116 _Int64Codes as _Int64Codes,
117 _Float16Codes as _Float16Codes,
118 _Float32Codes as _Float32Codes,
119 _Float64Codes as _Float64Codes,
120 _Complex64Codes as _Complex64Codes,
121 _Complex128Codes as _Complex128Codes,
122 _ByteCodes as _ByteCodes,
123 _ShortCodes as _ShortCodes,
124 _IntCCodes as _IntCCodes,
125 _IntPCodes as _IntPCodes,
126 _IntCodes as _IntCodes,
127 _LongLongCodes as _LongLongCodes,
128 _UByteCodes as _UByteCodes,
129 _UShortCodes as _UShortCodes,
130 _UIntCCodes as _UIntCCodes,
131 _UIntPCodes as _UIntPCodes,
132 _UIntCodes as _UIntCodes,
133 _ULongLongCodes as _ULongLongCodes,
134 _HalfCodes as _HalfCodes,
135 _SingleCodes as _SingleCodes,
136 _DoubleCodes as _DoubleCodes,
137 _LongDoubleCodes as _LongDoubleCodes,
138 _CSingleCodes as _CSingleCodes,
139 _CDoubleCodes as _CDoubleCodes,
140 _CLongDoubleCodes as _CLongDoubleCodes,
141 _DT64Codes as _DT64Codes,
142 _TD64Codes as _TD64Codes,
143 _StrCodes as _StrCodes,
144 _BytesCodes as _BytesCodes,
145 _VoidCodes as _VoidCodes,
146 _ObjectCodes as _ObjectCodes,
147)
148from ._scalars import (
149 _CharLike_co as _CharLike_co,
150 _BoolLike_co as _BoolLike_co,
151 _UIntLike_co as _UIntLike_co,
152 _IntLike_co as _IntLike_co,
153 _FloatLike_co as _FloatLike_co,
154 _ComplexLike_co as _ComplexLike_co,
155 _TD64Like_co as _TD64Like_co,
156 _NumberLike_co as _NumberLike_co,
157 _ScalarLike_co as _ScalarLike_co,
158 _VoidLike_co as _VoidLike_co,
159)
160from ._shape import (
161 _Shape as _Shape,
162 _ShapeLike as _ShapeLike,
163)
164from ._dtype_like import (
165 DTypeLike as DTypeLike,
166 _DTypeLike as _DTypeLike,
167 _SupportsDType as _SupportsDType,
168 _VoidDTypeLike as _VoidDTypeLike,
169 _DTypeLikeBool as _DTypeLikeBool,
170 _DTypeLikeUInt as _DTypeLikeUInt,
171 _DTypeLikeInt as _DTypeLikeInt,
172 _DTypeLikeFloat as _DTypeLikeFloat,
173 _DTypeLikeComplex as _DTypeLikeComplex,
174 _DTypeLikeTD64 as _DTypeLikeTD64,
175 _DTypeLikeDT64 as _DTypeLikeDT64,
176 _DTypeLikeObject as _DTypeLikeObject,
177 _DTypeLikeVoid as _DTypeLikeVoid,
178 _DTypeLikeStr as _DTypeLikeStr,
179 _DTypeLikeBytes as _DTypeLikeBytes,
180 _DTypeLikeComplex_co as _DTypeLikeComplex_co,
181)
182from ._array_like import (
183 ArrayLike as ArrayLike,
184 _ArrayLike as _ArrayLike,
185 _FiniteNestedSequence as _FiniteNestedSequence,
186 _SupportsArray as _SupportsArray,
187 _SupportsArrayFunc as _SupportsArrayFunc,
188 _ArrayLikeInt as _ArrayLikeInt,
189 _ArrayLikeBool_co as _ArrayLikeBool_co,
190 _ArrayLikeUInt_co as _ArrayLikeUInt_co,
191 _ArrayLikeInt_co as _ArrayLikeInt_co,
192 _ArrayLikeFloat_co as _ArrayLikeFloat_co,
193 _ArrayLikeComplex_co as _ArrayLikeComplex_co,
194 _ArrayLikeNumber_co as _ArrayLikeNumber_co,
195 _ArrayLikeTD64_co as _ArrayLikeTD64_co,
196 _ArrayLikeDT64_co as _ArrayLikeDT64_co,
197 _ArrayLikeObject_co as _ArrayLikeObject_co,
198 _ArrayLikeVoid_co as _ArrayLikeVoid_co,
199 _ArrayLikeStr_co as _ArrayLikeStr_co,
200 _ArrayLikeBytes_co as _ArrayLikeBytes_co,
201 _ArrayLikeUnknown as _ArrayLikeUnknown,
202 _UnknownType as _UnknownType,
203)
204from ._generic_alias import (
205 NDArray as NDArray,
206 _DType as _DType,
207 _GenericAlias as _GenericAlias,
208)
209
210if TYPE_CHECKING:
211 from ._ufunc import (
212 _UFunc_Nin1_Nout1 as _UFunc_Nin1_Nout1,
213 _UFunc_Nin2_Nout1 as _UFunc_Nin2_Nout1,
214 _UFunc_Nin1_Nout2 as _UFunc_Nin1_Nout2,
215 _UFunc_Nin2_Nout2 as _UFunc_Nin2_Nout2,
216 _GUFunc_Nin2_Nout1 as _GUFunc_Nin2_Nout1,
217 )
218else:
219 # Declare the (type-check-only) ufunc subclasses as ufunc aliases during
220 # runtime; this helps autocompletion tools such as Jedi (numpy/numpy#19834)
221 _UFunc_Nin1_Nout1 = ufunc
222 _UFunc_Nin2_Nout1 = ufunc
223 _UFunc_Nin1_Nout2 = ufunc
224 _UFunc_Nin2_Nout2 = ufunc
225 _GUFunc_Nin2_Nout1 = ufunc