diff --git a/xarray/namedarray/parallelcompat.py b/xarray/namedarray/parallelcompat.py index 8a68f5e9562..f3a65db5326 100644 --- a/xarray/namedarray/parallelcompat.py +++ b/xarray/namedarray/parallelcompat.py @@ -156,38 +156,43 @@ def get_chunked_array_type(*args: Any) -> ChunkManagerEntrypoint[Any]: if is_chunked_array(a) and type(a) not in ALLOWED_NON_CHUNKED_TYPES ] - # Asserts all arrays are the same type (or numpy etc.) - chunked_array_types = {type(a) for a in chunked_arrays} - if len(chunked_array_types) > 1: - raise TypeError( - f"Mixing chunked array types is not supported, but received multiple types: {chunked_array_types}" - ) - elif len(chunked_array_types) == 0: + if not chunked_arrays: raise TypeError("Expected a chunked array but none were found") - # iterate over defined chunk managers, seeing if each recognises this array type - chunked_arr = chunked_arrays[0] + # iterate over defined chunk managers, checking that every chunked array + # is recognised by the same manager (i.e. not a mix of cubed and dask). + # A manager may recognise more than one array class (e.g. dask and a + # subclass of dask.Array), so compare the selected managers rather than + # the exact array types (#11539). chunkmanagers = list_chunkmanagers() - selected = [ - chunkmanager - for chunkmanager in chunkmanagers.values() - if chunkmanager.is_chunked_array(chunked_arr) - ] - if not selected: - if ( - maybe_lib := type(chunked_arr).__module__.split(".")[0] - ) in KNOWN_CHUNKMANAGERS: - suggestion = f"Please try installing {KNOWN_CHUNKMANAGERS[maybe_lib]!r}." - else: - suggestion = "This is usually the result of a missing dependency." - raise TypeError( - f"Could not find a Chunk Manager which recognises type {type(chunked_arr)}" - f" {suggestion}" - ) - elif len(selected) >= 2: - raise TypeError(f"Multiple ChunkManagers recognise type {type(chunked_arr)}") - else: - return selected[0] + selected: list[ChunkManagerEntrypoint[Any]] | None = None + for a in chunked_arrays: + recognising = [ + chunkmanager + for chunkmanager in chunkmanagers.values() + if chunkmanager.is_chunked_array(a) + ] + if not recognising: + if (maybe_lib := type(a).__module__.split(".")[0]) in KNOWN_CHUNKMANAGERS: + suggestion = ( + f"Please try installing {KNOWN_CHUNKMANAGERS[maybe_lib]!r}." + ) + else: + suggestion = "This is usually the result of a missing dependency." + raise TypeError( + f"Could not find a Chunk Manager which recognises type {type(a)}" + f" {suggestion}" + ) + elif len(recognising) >= 2: + raise TypeError(f"Multiple ChunkManagers recognise type {type(a)}") + if selected is None: + selected = recognising + elif recognising[0] is not selected[0]: + raise TypeError( + "Mixing chunked array types is not supported, but received" + f" multiple types: {{ {', '.join(str(type(x)) for x in chunked_arrays)} }}" + ) + return selected[0] class ChunkManagerEntrypoint(ABC, Generic[T_ChunkedArray]): diff --git a/xarray/tests/test_parallelcompat.py b/xarray/tests/test_parallelcompat.py index 456c4cae46d..60f5c688b70 100644 --- a/xarray/tests/test_parallelcompat.py +++ b/xarray/tests/test_parallelcompat.py @@ -261,6 +261,22 @@ def test_raise_on_mixed_array_types(self, register_dummy_chunkmanager) -> None: with pytest.raises(TypeError, match="received multiple types"): get_chunked_array_type(*[dask_arr, dummy_arr]) + @requires_dask + def test_subclass_of_chunked_array_type(self, register_dummy_chunkmanager) -> None: + # Regression test for #11539: two arrays recognised by the same + # chunk manager (e.g. dask.Array and a subclass of it) must not be + # treated as a mix of chunked array types. + import dask.array as da + + class SubDaskArray(da.Array): + pass + + plain = da.from_array([1, 2, 3], chunks=(1,)) + subclass = SubDaskArray(plain.dask, plain.name, plain.chunks, plain.dtype) + + manager = get_chunked_array_type(plain, subclass) + assert isinstance(manager, DaskManager) + def test_bogus_entrypoint() -> None: # Create a bogus entry-point as if the user broke their setup.cfg