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