__init__.py 6.9 KB

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