From 7d3fde0e075001da44cccfd44826b6b7187197a1 Mon Sep 17 00:00:00 2001 From: IMGillusion Date: Sun, 6 Sep 2026 03:25:19 +0800 Subject: [PATCH 1/2] Fix get_chunked_array_type rejecting arrays of the same chunk manager Compare the chunk managers that recognise each array instead of the exact array types, so a dask array and a subclass of dask.Array are no longer reported as a mix of chunked array types. A real mix of frameworks (e.g. dask and cubed) still raises TypeError. Fixes #11539 --- xarray/namedarray/parallelcompat.py | 63 ++++++++++++++++------------- xarray/tests/test_parallelcompat.py | 16 ++++++++ 2 files changed, 50 insertions(+), 29 deletions(-) diff --git a/xarray/namedarray/parallelcompat.py b/xarray/namedarray/parallelcompat.py index 8a68f5e9562..8a8cb4e998d 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 From da143cfd75b09648feb18d3034eefd63701ebd13 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 5 Sep 2026 19:25:55 +0000 Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- xarray/namedarray/parallelcompat.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/xarray/namedarray/parallelcompat.py b/xarray/namedarray/parallelcompat.py index 8a8cb4e998d..f3a65db5326 100644 --- a/xarray/namedarray/parallelcompat.py +++ b/xarray/namedarray/parallelcompat.py @@ -173,10 +173,10 @@ def get_chunked_array_type(*args: Any) -> ChunkManagerEntrypoint[Any]: 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}." + 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(