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

Code before:
    -----
    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.

        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
            If True, the graph is optimized before rendering.  Otherwise,

Review comment (performance):
This PR breaks pint, because just like xarray it violates the dask collection protocol by defining a `__dask_graph__()` method that may return None:

```python
>>> import dask.array as da, pint
>>> q = pint.Quantity(1, "m")
>>> q.__dask_graph__()
(None)
>>> q = pint.Quantity(da.zeros(5), "m")
>>> q.__dask_graph__()
HighLevelGraph with 1 layers.
<dask.highlevelgraph.HighLevelGraph object at 0x7f043c3889a0>
 0. zeros_like-902bf3bcabaabb3fedce3f347fb26081
```

There may be more third-party libraries in the wild that do the same.
A blocklist approach would be more robust:

```python
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]
# Add here other third-party packages where you don't want
# to call the `__dask_graph__` method - typically because it's expensive
if pkg_name == "dask_expr":
    return True

# xarray, pint, and possibly other libraries return None when they wrap a non-dask object
return x.__dask_graph__() is not None
```

Code after:
    -----
    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
        ----------
        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
            If True, the graph is optimized before rendering.  Otherwise,
            the graph is displayed as is. Default is False.