_array_api_info.pyi 4.7 KB

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