Another mock duck array class (like EncapsulateNDArray), but designed to be above Dask in the type casting hierarchy (that is, WrappedArray wraps Dask Array) and be even more minimal in API. Tests that Dask defers properly to upcast types.
| 97 | |
| 98 | |
| 99 | class WrappedArray(np.lib.mixins.NDArrayOperatorsMixin): |
| 100 | """ |
| 101 | Another mock duck array class (like EncapsulateNDArray), but |
| 102 | designed to be above Dask in the type casting hierarchy (that is, |
| 103 | WrappedArray wraps Dask Array) and be even more minimal in API. |
| 104 | Tests that Dask defers properly to upcast types. |
| 105 | """ |
| 106 | |
| 107 | def __init__(self, arr, **attrs): |
| 108 | self.arr = arr |
| 109 | self.attrs = attrs |
| 110 | |
| 111 | def __array__(self, *args, **kwargs): |
| 112 | return np.asarray(self.arr, *args, **kwargs) |
| 113 | |
| 114 | def _downcast_args(self, args): |
| 115 | for arg in args: |
| 116 | if isinstance(arg, type(self)): |
| 117 | yield arg.arr |
| 118 | elif isinstance(arg, (tuple, list)): |
| 119 | yield type(arg)(self._downcast_args(arg)) |
| 120 | else: |
| 121 | yield arg |
| 122 | |
| 123 | def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): |
| 124 | inputs = tuple(self._downcast_args(inputs)) |
| 125 | return type(self)(getattr(ufunc, method)(*inputs, **kwargs), **self.attrs) |
| 126 | |
| 127 | def __array_function__(self, func, types, args, kwargs): |
| 128 | args = tuple(self._downcast_args(args)) |
| 129 | return type(self)(func(*args, **kwargs), **self.attrs) |
| 130 | |
| 131 | def __dask_graph__(self): |
| 132 | # Note: make sure that dask dusk arrays do not interfere with the |
| 133 | # dispatch mechanism. The return value here, doesn't matter. |
| 134 | return ... |
| 135 | |
| 136 | shape = dispatch_property("shape") |
| 137 | ndim = dispatch_property("ndim") |
| 138 | dtype = dispatch_property("dtype") |
| 139 | |
| 140 | def __getitem__(self, key): |
| 141 | return type(self)(self.arr[key], **self.attrs) |
| 142 | |
| 143 | def __setitem__(self, key, value): |
| 144 | self.arr[key] = value |
| 145 | |
| 146 | |
| 147 | @pytest.mark.parametrize( |