array-api-typing/tests/integration/test_numpy2p0.pyi at b7fa49dfd28f7f84f44d842fc4c8832ba0d0eb28 · data-apis/array-api-typing · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
# mypy: disable-error-code="no-redef, explicit-any"

from types import ModuleType
from typing import Any, TypeAlias

import numpy as np
import numpy.typing as npt

import array_api_typing as xpt

# DType aliases
F32: TypeAlias = np.float32
I32: TypeAlias = np.int32

# Define NDArrays against which we can test the protocols
nparr: npt.NDArray[Any]
nparr_i32: npt.NDArray[I32]
nparr_f32: npt.NDArray[F32]
nparr_b: npt.NDArray[np.bool_]

# =========================================================
# `xpt.HasArrayNamespace`

# Check assignment
_001: xpt.HasArrayNamespace[ModuleType] = nparr
_002: xpt.HasArrayNamespace[ModuleType] = nparr_i32
_003: xpt.HasArrayNamespace[ModuleType] = nparr_f32
_004: xpt.HasArrayNamespace[ModuleType] = nparr_b

# Check `__array_namespace__` method
a_ns: xpt.HasArrayNamespace[ModuleType] = nparr
ns: ModuleType = a_ns.__array_namespace__()

# =========================================================
# `xpt.HasDType`

# Check DTypeT_co assignment
_005: xpt.HasDType[Any] = nparr
_006: xpt.HasDType[np.dtype[I32]] = nparr_i32
_007: xpt.HasDType[np.dtype[F32]] = nparr_f32
_008: xpt.HasDType[np.dtype[np.bool_]] = nparr_b

# =========================================================
# `xpt.Array`

# Check NamespaceT_co assignment
x_ns: xpt.Array[Any, ModuleType] = nparr

# Check DTypeT_co assignment
_009: xpt.Array[Any] = nparr
_010: xpt.Array[np.dtype[I32]] = nparr_i32
_011: xpt.Array[np.dtype[F32]] = nparr_f32
_012: xpt.Array[np.dtype[np.bool_]] = nparr_b