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

Code before:
def is_dask_collection(x) -> bool:
    """Returns ``True`` if ``x`` is a dask collection.

    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__:
            return x.__dask_graph__() is not None
        else:
            return hasattr(x, "__dask_graph__")

    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.

        Parameters
        ----------
        filename : str or None, optional
            The name of the file to write to disk. If the provided `filename`
            doesn't include an extension, '.png' will be used by default.
            If `filename` is None, no file will be written, and we communicate
            with dot using only pipes.
        format : {'png', 'pdf', 'dot', 'svg', 'jpeg', 'jpg'}, optional
            Format in which to write output file.  Default is 'png'.
        optimize_graph : bool, optional

Review comment (performance):
@phofl this change is actually kinda important.  We were generating the task graph before whenever calling `is_dask_collection`.  Maybe caching was saving us here, but this could be unpleasantly expensive for dask-expr.

The history here is that xarray collections have a `__dask_graph__` attribute, but are not always dask collections (they sometimes hold dask arrays but sometimes just have numpy arrays).

Code after:
def is_dask_collection(x) -> bool:
    """Returns ``True`` if ``x`` is a dask collection.

    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.