A class that "mocks" ndarray by encapsulating an ndarray and using protocols to "look like" an ndarray. Basically tests whether Dask works fine with something that is essentially an array but uses protocols instead of being an actual array. Must be manually registered as a valid
| 39 | |
| 40 | |
| 41 | class EncapsulateNDArray(np.lib.mixins.NDArrayOperatorsMixin): |
| 42 | """ |
| 43 | A class that "mocks" ndarray by encapsulating an ndarray and using |
| 44 | protocols to "look like" an ndarray. Basically tests whether Dask |
| 45 | works fine with something that is essentially an array but uses |
| 46 | protocols instead of being an actual array. Must be manually |
| 47 | registered as a valid chunk type to be considered a downcast type |
| 48 | of Dask array in the type casting hierarchy. |
| 49 | """ |
| 50 | |
| 51 | __array_priority__ = 20 |
| 52 | |
| 53 | def __init__(self, arr): |
| 54 | self.arr = arr |
| 55 | |
| 56 | def __array__(self, *args, **kwargs): |
| 57 | return np.asarray(self.arr, *args, **kwargs) |
| 58 | |
| 59 | def __array_function__(self, f, t, arrs, kw): |
| 60 | if not all( |
| 61 | issubclass(ti, (type(self), np.ndarray) + np.ScalarType) for ti in t |
| 62 | ): |
| 63 | return NotImplemented |
| 64 | arrs = tuple( |
| 65 | arr if not isinstance(arr, type(self)) else arr.arr for arr in arrs |
| 66 | ) |
| 67 | t = tuple(ti for ti in t if not issubclass(ti, type(self))) |
| 68 | print(t) |
| 69 | a = self.arr.__array_function__(f, t, arrs, kw) |
| 70 | return a if not isinstance(a, np.ndarray) else type(self)(a) |
| 71 | |
| 72 | __getitem__ = wrap("__getitem__") |
| 73 | |
| 74 | __setitem__ = wrap("__setitem__") |
| 75 | |
| 76 | def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): |
| 77 | if not all( |
| 78 | isinstance(i, (type(self), np.ndarray) + np.ScalarType) for i in inputs |
| 79 | ): |
| 80 | return NotImplemented |
| 81 | inputs = tuple(i if not isinstance(i, type(self)) else i.arr for i in inputs) |
| 82 | a = getattr(ufunc, method)(*inputs, **kwargs) |
| 83 | return a if not isinstance(a, np.ndarray) else type(self)(a) |
| 84 | |
| 85 | shape = dispatch_property("shape") |
| 86 | ndim = dispatch_property("ndim") |
| 87 | dtype = dispatch_property("dtype") |
| 88 | |
| 89 | astype = wrap("astype") |
| 90 | sum = wrap("sum") |
| 91 | prod = wrap("prod") |
| 92 | reshape = wrap("reshape") |
| 93 | squeeze = wrap("squeeze") |
| 94 | |
| 95 | |
| 96 | da.register_chunk_type(EncapsulateNDArray) |