1
1
from __future__ import annotations
2
2
3
3
from importlib import import_module
4
- from typing import Any , Literal
4
+ from types import ModuleType
5
+ from typing import TYPE_CHECKING , Any , Literal , Type
5
6
6
7
import numpy as np
7
8
from packaging .version import Version
10
11
11
12
integer_types = (int , np .integer )
12
13
13
- ModType = Literal ["dask" , "pint" , "cupy" , "sparse" ]
14
+ if TYPE_CHECKING :
15
+ ModType = Literal ["dask" , "pint" , "cupy" , "sparse" ]
16
+ DuckArrayTypes = tuple [Type [Any ], ...] # TODO: improve this? maybe Generic
14
17
15
18
16
19
class DuckArrayModule :
@@ -21,12 +24,15 @@ class DuckArrayModule:
21
24
https://github.com/pydata/xarray/pull/5561#discussion_r664815718
22
25
"""
23
26
24
- module : ModType | None
27
+ module : ModuleType | None
25
28
version : Version
26
- type : tuple [ type [ Any ]] # TODO: improve this? maybe Generic
29
+ type : DuckArrayTypes
27
30
available : bool
28
31
29
32
def __init__ (self , mod : ModType ) -> None :
33
+ duck_array_module : ModuleType | None = None
34
+ duck_array_version : Version
35
+ duck_array_type : DuckArrayTypes
30
36
try :
31
37
duck_array_module = import_module (mod )
32
38
duck_array_version = Version (duck_array_module .__version__ )
@@ -53,7 +59,7 @@ def __init__(self, mod: ModType) -> None:
53
59
self .available = duck_array_module is not None
54
60
55
61
56
- def array_type (mod : ModType ) -> tuple [ type [ Any ]] :
62
+ def array_type (mod : ModType ) -> DuckArrayTypes :
57
63
"""Quick wrapper to get the array class of the module."""
58
64
return DuckArrayModule (mod ).type
59
65
0 commit comments