Fix get_chunked_array_type rejecting arrays of the same chunk manager

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]):

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