|
| 1 | +# mypy: disable-error-code="no-redef, explicit-any" |
| 2 | + |
| 3 | +from types import ModuleType |
| 4 | +from typing import TypeAlias |
| 5 | + |
| 6 | +import dask.array as da |
| 7 | +import numpy as np |
| 8 | + |
| 9 | +import array_api_typing as xpt |
| 10 | + |
| 11 | +# DType aliases |
| 12 | +F32: TypeAlias = np.float32 |
| 13 | +I32: TypeAlias = np.int32 |
| 14 | + |
| 15 | +# Define Dask Arrays against which we can test the protocols |
| 16 | +# (Dask's Array type is not parametrized by dtype at type level.) |
| 17 | +darr: da.Array |
| 18 | +darr_i32: da.Array |
| 19 | +darr_f32: da.Array |
| 20 | +darr_b: da.Array |
| 21 | + |
| 22 | +# ========================================================= |
| 23 | +# `xpt.HasArrayNamespace` |
| 24 | + |
| 25 | +# Check assignment |
| 26 | +_001: xpt.HasArrayNamespace[ModuleType] = darr |
| 27 | +_002: xpt.HasArrayNamespace[ModuleType] = darr_i32 |
| 28 | +_003: xpt.HasArrayNamespace[ModuleType] = darr_f32 |
| 29 | +_004: xpt.HasArrayNamespace[ModuleType] = darr_b |
| 30 | + |
| 31 | +# Check `__array_namespace__` method |
| 32 | +a_ns: xpt.HasArrayNamespace[ModuleType] = darr |
| 33 | +ns: ModuleType = a_ns.__array_namespace__() |
| 34 | + |
| 35 | +# ========================================================= |
| 36 | +# `xpt.HasDType` |
| 37 | + |
| 38 | +# Check DTypeT_co assignment |
| 39 | +_005: xpt.HasDType[Any] = darr |
| 40 | +_006: xpt.HasDType[np.dtype[I32]] = darr_i32 |
| 41 | +_007: xpt.HasDType[np.dtype[F32]] = darr_f32 |
| 42 | +_008: xpt.HasDType[np.dtype[np.bool_]] = darr_b |
| 43 | + |
| 44 | +# ========================================================= |
| 45 | +# `xpt.Array` |
| 46 | + |
| 47 | +# Check NamespaceT_co assignment |
| 48 | +x_ns: xpt.Array[Any, ModuleType] = darr |
| 49 | + |
| 50 | +# Check DTypeT_co assignment |
| 51 | +_009: xpt.Array[Any] = darr |
| 52 | +_010: xpt.Array[np.dtype[I32]] = darr_i32 |
| 53 | +_011: xpt.Array[np.dtype[F32]] = darr_f32 |
| 54 | +_012: xpt.Array[np.dtype[np.bool_]] = darr_b |
0 commit comments