Repository: dask/dask
File: dask/base.py (Python)

Code before:
    Parameters
    ----------
    x : Any
        Object to test.

    Returns
    -------
    result : bool
        ``True`` if `x` is a Dask collection.

    Notes
    -----
    The DaskCollection typing.Protocol implementation defines a Dask
    collection as a class that returns a Mapping from the
    ``__dask_graph__`` method. This helper function existed before the
    implementation of the protocol.

    """
    try:
        if "xarray" in type(x).__module__:
            import xarray

            if isinstance(x, xarray.Dataset):
                return any(is_dask_collection(v._data) for _, v in x.variables.items())
            elif isinstance(x, xarray.DataArray):
                return is_dask_collection(x.variable._data)
            elif isinstance(x, xarray.Variable):
                return is_dask_collection(x._data)
            else:
                raise TypeError("Unfamiliar with xarray type", type(x))
        else:
            return (
                hasattr(x, "__dask_graph__")
                and callable(x.__dask_graph__)
                and not isinstance(x, type)
            )

    except (AttributeError, TypeError):
        return False


class DaskMethodsMixin:
    """A mixin adding standard dask collection methods"""

    __slots__ = ()

    def visualize(self, filename="mydask", format=None, optimize_graph=False, **kwargs):
        """Render the computation of this object's task graph using graphviz.

        Requires ``graphviz`` to be installed.


Review comment (suggestion):
```suggestion
            elif isinstance(x, xarray.DataArray):
                return is_dask_collection(x.variable._data) or any(is_dask_collection(var._data) for var in x._coords.values())
```

Code after:
    Parameters
    ----------
    x : Any
        Object to test.

    Returns
    -------
    result : bool
        ``True`` if `x` is a Dask collection.

    Notes
    -----
    The DaskCollection typing.Protocol implementation defines a Dask
    collection as a class that returns a Mapping from the
    ``__dask_graph__`` method. This helper function existed before the
    implementation of the protocol.

    """
    if (
        isinstance(x, type)
        or not hasattr(x, "__dask_graph__")
        or not callable(x.__dask_graph__)
    ):
        return False

    pkg_name = getattr(type(x), "__module__", "").split(".")[0]
    if pkg_name == "dask_expr":
        # Temporary hack to avoid graph materialization. Note that this won't work with
        # dask_expr.array objects wrapped by xarray or pint. By the time dask_expr.array
        # is published, we hope to be able to rewrite this method completely.
        # Read: https://github.com/dask/dask/pull/10676
        return True

    # xarray, pint, and possibly other wrappers always define a __dask_graph__ method,
    # but it may return None if they wrap around a non-dask object.
    # In all known dask collections other than dask-expr,
    # calling __dask_graph__ is cheap.
    return x.__dask_graph__() is not None


class DaskMethodsMixin:
    """A mixin adding standard dask collection methods"""

    __slots__ = ()

    def visualize(self, filename="mydask", format=None, optimize_graph=False, **kwargs):
        """Render the computation of this object's task graph using graphviz.

        Requires ``graphviz`` to be installed.

        Parameters