1from distutils.version import LooseVersion
2from importlib import import_module
3
4import numpy as np
5
6from .utils import is_duck_array
7
8integer_types = (int, np.integer)
9
10
11class DuckArrayModule:
12    """
13    Solely for internal isinstance and version checks.
14
15    Motivated by having to only import pint when required (as pint currently imports xarray)
16    https://github.com/pydata/xarray/pull/5561#discussion_r664815718
17    """
18
19    def __init__(self, mod):
20        try:
21            duck_array_module = import_module(mod)
22            duck_array_version = LooseVersion(duck_array_module.__version__)
23
24            if mod == "dask":
25                duck_array_type = (import_module("dask.array").Array,)
26            elif mod == "pint":
27                duck_array_type = (duck_array_module.Quantity,)
28            elif mod == "cupy":
29                duck_array_type = (duck_array_module.ndarray,)
30            elif mod == "sparse":
31                duck_array_type = (duck_array_module.SparseArray,)
32            else:
33                raise NotImplementedError
34
35        except ImportError:  # pragma: no cover
36            duck_array_module = None
37            duck_array_version = LooseVersion("0.0.0")
38            duck_array_type = ()
39
40        self.module = duck_array_module
41        self.version = duck_array_version
42        self.type = duck_array_type
43        self.available = duck_array_module is not None
44
45
46def is_duck_dask_array(x):
47    if DuckArrayModule("dask").available:
48        from dask.base import is_dask_collection
49
50        return is_duck_array(x) and is_dask_collection(x)
51    else:
52        return False
53
54
55dsk = DuckArrayModule("dask")
56dask_version = dsk.version
57dask_array_type = dsk.type
58
59sp = DuckArrayModule("sparse")
60sparse_array_type = sp.type
61sparse_version = sp.version
62
63cupy_array_type = DuckArrayModule("cupy").type
64