Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.9/dist-packages/xarray/namedarray/_typing.py: 22%

1from __future__ import annotations 

2 

3import sys 

4from collections.abc import Hashable, Iterable, Mapping, Sequence 

5from enum import Enum 

6from types import ModuleType 

7from typing import ( 

8 TYPE_CHECKING, 

9 Any, 

10 Callable, 

11 Final, 

12 Literal, 

13 Protocol, 

14 SupportsIndex, 

15 TypeVar, 

16 Union, 

17 overload, 

18 runtime_checkable, 

19) 

20 

21import numpy as np 

22 

23try: 

24 if sys.version_info >= (3, 11): 

25 from typing import TypeAlias 

26 else: 

27 from typing_extensions import TypeAlias 

28except ImportError: 

29 if TYPE_CHECKING: 

30 raise 

31 else: 

32 Self: Any = None 

33 

34 

35# Singleton type, as per https://github.com/python/typing/pull/240 

36class Default(Enum): 

37 token: Final = 0 

38 

39 

40_default = Default.token 

41 

42# https://stackoverflow.com/questions/74633074/how-to-type-hint-a-generic-numpy-array 

43_T = TypeVar("_T") 

44_T_co = TypeVar("_T_co", covariant=True) 

45 

46_dtype = np.dtype 

47_DType = TypeVar("_DType", bound=np.dtype[Any]) 

48_DType_co = TypeVar("_DType_co", covariant=True, bound=np.dtype[Any]) 

49# A subset of `npt.DTypeLike` that can be parametrized w.r.t. `np.generic` 

50 

51_ScalarType = TypeVar("_ScalarType", bound=np.generic) 

52_ScalarType_co = TypeVar("_ScalarType_co", bound=np.generic, covariant=True) 

53 

54 

55# A protocol for anything with the dtype attribute 

56@runtime_checkable 

57class _SupportsDType(Protocol[_DType_co]): 

58 @property 

59 def dtype(self) -> _DType_co: ... 

60 

61 

62_DTypeLike = Union[ 

63 np.dtype[_ScalarType], 

64 type[_ScalarType], 

65 _SupportsDType[np.dtype[_ScalarType]], 

66] 

67 

68# For unknown shapes Dask uses np.nan, array_api uses None: 

69_IntOrUnknown = int 

70_Shape = tuple[_IntOrUnknown, ...] 

71_ShapeLike = Union[SupportsIndex, Sequence[SupportsIndex]] 

72_ShapeType = TypeVar("_ShapeType", bound=Any) 

73_ShapeType_co = TypeVar("_ShapeType_co", bound=Any, covariant=True) 

74 

75_Axis = int 

76_Axes = tuple[_Axis, ...] 

77_AxisLike = Union[_Axis, _Axes] 

78 

79_Chunks = tuple[_Shape, ...] 

80_NormalizedChunks = tuple[tuple[int, ...], ...] 

81# FYI in some cases we don't allow `None`, which this doesn't take account of. 

82T_ChunkDim: TypeAlias = Union[int, Literal["auto"], None, tuple[int, ...]] 

83# We allow the tuple form of this (though arguably we could transition to named dims only) 

84T_Chunks: TypeAlias = Union[T_ChunkDim, Mapping[Any, T_ChunkDim]] 

85 

86_Dim = Hashable 

87_Dims = tuple[_Dim, ...] 

88 

89_DimsLike = Union[str, Iterable[_Dim]] 

90 

91# https://data-apis.org/array-api/latest/API_specification/indexing.html 

92# TODO: np.array_api was bugged and didn't allow (None,), but should! 

93# https://github.com/numpy/numpy/pull/25022 

94# https://github.com/data-apis/array-api/pull/674 

95_IndexKey = Union[int, slice, "ellipsis"] 

96_IndexKeys = tuple[Union[_IndexKey], ...] # tuple[Union[_IndexKey, None], ...] 

97_IndexKeyLike = Union[_IndexKey, _IndexKeys] 

98 

99_AttrsLike = Union[Mapping[Any, Any], None] 

100 

101 

102class _SupportsReal(Protocol[_T_co]): 

103 @property 

104 def real(self) -> _T_co: ... 

105 

106 

107class _SupportsImag(Protocol[_T_co]): 

108 @property 

109 def imag(self) -> _T_co: ... 

110 

111 

112@runtime_checkable 

113class _array(Protocol[_ShapeType_co, _DType_co]): 

114 """ 

115 Minimal duck array named array uses. 

116 

117 Corresponds to np.ndarray. 

118 """ 

119 

120 @property 

121 def shape(self) -> _Shape: ... 

122 

123 @property 

124 def dtype(self) -> _DType_co: ... 

125 

126 

127@runtime_checkable 

128class _arrayfunction( 

129 _array[_ShapeType_co, _DType_co], Protocol[_ShapeType_co, _DType_co] 

130): 

131 """ 

132 Duck array supporting NEP 18. 

133 

134 Corresponds to np.ndarray. 

135 """ 

136 

137 @overload 

138 def __getitem__( 

139 self, key: _arrayfunction[Any, Any] | tuple[_arrayfunction[Any, Any], ...], / 

140 ) -> _arrayfunction[Any, _DType_co]: ... 

141 

142 @overload 

143 def __getitem__(self, key: _IndexKeyLike, /) -> Any: ... 

144 

145 def __getitem__( 

146 self, 

147 key: ( 

148 _IndexKeyLike 

149 | _arrayfunction[Any, Any] 

150 | tuple[_arrayfunction[Any, Any], ...] 

151 ), 

152 /, 

153 ) -> _arrayfunction[Any, _DType_co] | Any: ... 

154 

155 @overload 

156 def __array__( 

157 self, dtype: None = ..., /, *, copy: None | bool = ... 

158 ) -> np.ndarray[Any, _DType_co]: ... 

159 @overload 

160 def __array__( 

161 self, dtype: _DType, /, *, copy: None | bool = ... 

162 ) -> np.ndarray[Any, _DType]: ... 

163 

164 def __array__( 

165 self, dtype: _DType | None = ..., /, *, copy: None | bool = ... 

166 ) -> np.ndarray[Any, _DType] | np.ndarray[Any, _DType_co]: ... 

167 

168 # TODO: Should return the same subclass but with a new dtype generic. 

169 # https://github.com/python/typing/issues/548 

170 def __array_ufunc__( 

171 self, 

172 ufunc: Any, 

173 method: Any, 

174 *inputs: Any, 

175 **kwargs: Any, 

176 ) -> Any: ... 

177 

178 # TODO: Should return the same subclass but with a new dtype generic. 

179 # https://github.com/python/typing/issues/548 

180 def __array_function__( 

181 self, 

182 func: Callable[..., Any], 

183 types: Iterable[type], 

184 args: Iterable[Any], 

185 kwargs: Mapping[str, Any], 

186 ) -> Any: ... 

187 

188 @property 

189 def imag(self) -> _arrayfunction[_ShapeType_co, Any]: ... 

190 

191 @property 

192 def real(self) -> _arrayfunction[_ShapeType_co, Any]: ... 

193 

194 

195@runtime_checkable 

196class _arrayapi(_array[_ShapeType_co, _DType_co], Protocol[_ShapeType_co, _DType_co]): 

197 """ 

198 Duck array supporting NEP 47. 

199 

200 Corresponds to np.ndarray. 

201 """ 

202 

203 def __getitem__( 

204 self, 

205 key: ( 

206 _IndexKeyLike | Any 

207 ), # TODO: Any should be _arrayapi[Any, _dtype[np.integer]] 

208 /, 

209 ) -> _arrayapi[Any, Any]: ... 

210 

211 def __array_namespace__(self) -> ModuleType: ... 

212 

213 

214# NamedArray can most likely use both __array_function__ and __array_namespace__: 

215_arrayfunction_or_api = (_arrayfunction, _arrayapi) 

216 

217duckarray = Union[ 

218 _arrayfunction[_ShapeType_co, _DType_co], _arrayapi[_ShapeType_co, _DType_co] 

219] 

220 

221# Corresponds to np.typing.NDArray: 

222DuckArray = _arrayfunction[Any, np.dtype[_ScalarType_co]] 

223 

224 

225@runtime_checkable 

226class _chunkedarray( 

227 _array[_ShapeType_co, _DType_co], Protocol[_ShapeType_co, _DType_co] 

228): 

229 """ 

230 Minimal chunked duck array. 

231 

232 Corresponds to np.ndarray. 

233 """ 

234 

235 @property 

236 def chunks(self) -> _Chunks: ... 

237 

238 

239@runtime_checkable 

240class _chunkedarrayfunction( 

241 _arrayfunction[_ShapeType_co, _DType_co], Protocol[_ShapeType_co, _DType_co] 

242): 

243 """ 

244 Chunked duck array supporting NEP 18. 

245 

246 Corresponds to np.ndarray. 

247 """ 

248 

249 @property 

250 def chunks(self) -> _Chunks: ... 

251 

252 

253@runtime_checkable 

254class _chunkedarrayapi( 

255 _arrayapi[_ShapeType_co, _DType_co], Protocol[_ShapeType_co, _DType_co] 

256): 

257 """ 

258 Chunked duck array supporting NEP 47. 

259 

260 Corresponds to np.ndarray. 

261 """ 

262 

263 @property 

264 def chunks(self) -> _Chunks: ... 

265 

266 

267# NamedArray can most likely use both __array_function__ and __array_namespace__: 

268_chunkedarrayfunction_or_api = (_chunkedarrayfunction, _chunkedarrayapi) 

269chunkedduckarray = Union[ 

270 _chunkedarrayfunction[_ShapeType_co, _DType_co], 

271 _chunkedarrayapi[_ShapeType_co, _DType_co], 

272] 

273 

274 

275@runtime_checkable 

276class _sparsearray( 

277 _array[_ShapeType_co, _DType_co], Protocol[_ShapeType_co, _DType_co] 

278): 

279 """ 

280 Minimal sparse duck array. 

281 

282 Corresponds to np.ndarray. 

283 """ 

284 

285 def todense(self) -> np.ndarray[Any, _DType_co]: ... 

286 

287 

288@runtime_checkable 

289class _sparsearrayfunction( 

290 _arrayfunction[_ShapeType_co, _DType_co], Protocol[_ShapeType_co, _DType_co] 

291): 

292 """ 

293 Sparse duck array supporting NEP 18. 

294 

295 Corresponds to np.ndarray. 

296 """ 

297 

298 def todense(self) -> np.ndarray[Any, _DType_co]: ... 

299 

300 

301@runtime_checkable 

302class _sparsearrayapi( 

303 _arrayapi[_ShapeType_co, _DType_co], Protocol[_ShapeType_co, _DType_co] 

304): 

305 """ 

306 Sparse duck array supporting NEP 47. 

307 

308 Corresponds to np.ndarray. 

309 """ 

310 

311 def todense(self) -> np.ndarray[Any, _DType_co]: ... 

312 

313 

314# NamedArray can most likely use both __array_function__ and __array_namespace__: 

315_sparsearrayfunction_or_api = (_sparsearrayfunction, _sparsearrayapi) 

316sparseduckarray = Union[ 

317 _sparsearrayfunction[_ShapeType_co, _DType_co], 

318 _sparsearrayapi[_ShapeType_co, _DType_co], 

319] 

320 

321ErrorOptions = Literal["raise", "ignore"] 

322ErrorOptionsWithWarn = Literal["raise", "warn", "ignore"]