<commit_before>from collections import OrderedDict
from collections.abc import Mapping, Iterator
from contextlib import contextmanager
from functools import partial
from hashlib import md5
from operator import getitem
import inspect
import pickle
import os
import threading
import uuid
from distutils.version import LooseVersion

from tlz import merge, groupby, curry, identity
from tlz.functoolz import Compose

from .compatibility import is_dataclass, dataclass_fields
from .context import thread_state
from .core import flatten, quote, get as simple_get, literal
from .hashing import hash_buffer_hex
from .utils import Dispatch, ensure_dict, apply
from . import config, local, threaded


__all__ = (
    "DaskMethodsMixin",
    "annotate",
    "is_dask_collection",
    "compute",
    "persist",
    "optimize",
    "visualize",
    "tokenize",
    "normalize_token",
)


@contextmanager
def annotate(**annotations):
    """Content Manager for setting HighLevelGraph Layer annotations.

    Annotations are metadata or soft constraints associated with
    tasks that dask schedulers may choose to respect: They signal intent
    without enforcing hard constraints. As such, they are
    primarily designed for use with the distributed scheduler.

    Almost any object can serve as an annotation, but small Python objects
    are preferred, while large objects such as NumPy arrays should be discouraged.

    Callables supplied as an annotation should take a single *key* argument and
    produce the appropriate annotation. Individual task keys in the annotated collection
    are supplied to the callable.

    Parameters
    ----------
    **annotations : key-value pairs

    Examples
    --------

    All tasks within array A should have priority 100 and be retried 3 times
    on failure.

    >>> with dask.annotate(priority=100, retries=3):
            A = da.ones((10000, 10000))

    Prioritise tasks within Array A on flattened block ID.

    >>> nblocks = (10, 10)
    >>> with dask.annotate(priority=lambda k: k[1]*nblocks[1] + k[2]):
            A = da.ones((1000, 1000), chunks=(100, 100))
    """

    prev_annotations = config.get("annotations", {})
    new_annotations = {
        **prev_annotations,
        **{f"annotations.{k}": v for k, v in annotations.items()},
    }

    with config.set(new_annotations):
        yield


def is_dask_collection(x):
    """Returns ``True`` if ``x`` is a dask collection"""
    try:
        return x.__dask_graph__() is not None
    except (AttributeError, TypeError):
        return False


class DaskMethodsMixin(object):
    """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.
<commit_msg>Layer annotation docstrings improvements (#6806)<commit_after>from collections import OrderedDict
from collections.abc import Mapping, Iterator
from contextlib import contextmanager
from functools import partial
from hashlib import md5
from operator import getitem
import inspect
import pickle
import os
import threading
import uuid
from distutils.version import LooseVersion

from tlz import merge, groupby, curry, identity
from tlz.functoolz import Compose

from .compatibility import is_dataclass, dataclass_fields
from .context import thread_state
from .core import flatten, quote, get as simple_get, literal
from .hashing import hash_buffer_hex
from .utils import Dispatch, ensure_dict, apply
from . import config, local, threaded


__all__ = (
    "DaskMethodsMixin",
    "annotate",
    "is_dask_collection",
    "compute",
    "persist",
    "optimize",
    "visualize",
    "tokenize",
    "normalize_token",
)


@contextmanager
def annotate(**annotations):
    """Context Manager for setting HighLevelGraph Layer annotations.

    Annotations are metadata or soft constraints associated with
    tasks that dask schedulers may choose to respect: They signal intent
    without enforcing hard constraints. As such, they are
    primarily designed for use with the distributed scheduler.

    Almost any object can serve as an annotation, but small Python objects
    are preferred, while large objects such as NumPy arrays are discouraged.

    Callables supplied as an annotation should take a single *key* argument and
    produce the appropriate annotation. Individual task keys in the annotated collection
    are supplied to the callable.

    Parameters
    ----------
    **annotations : key-value pairs

    Examples
    --------

    All tasks within array A should have priority 100 and be retried 3 times
    on failure.

    >>> import dask
    >>> import dask.array as da
    >>> with dask.annotate(priority=100, retries=3):
    ...     A = da.ones((10000, 10000))

    Prioritise tasks within Array A on flattened block ID.

    >>> nblocks = (10, 10)
    >>> with dask.annotate(priority=lambda k: k[1]*nblocks[1] + k[2]):
    ...     A = da.ones((1000, 1000), chunks=(100, 100))

    Annotations may be nested.

    >>> with dask.annotate(priority=1):
    ...     with dask.annotate(retries=3):
    ...         A = da.ones((1000, 1000))
    ...     B = A + 1
    """

    prev_annotations = config.get("annotations", {})
    new_annotations = {
        **prev_annotations,
        **{f"annotations.{k}": v for k, v in annotations.items()},
    }

    with config.set(new_annotations):
        yield


def is_dask_collection(x):
    """Returns ``True`` if ``x`` is a dask collection"""
    try:
        return x.__dask_graph__() is not None
    except (AttributeError, TypeError):
        return False


class DaskMethodsMixin(object):
    """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.

{'fraction_non_alphanumeric': 0.09043536728861994, 'fraction_numerical': 0.015969070432005378, 'mean_word_length': 3.4172234595397177, 'pattern_counts': {'":': 2, '<': 3, '<?xml version=': 0, '>': 30, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6145997', 'n_tokens_mistral': 1751, 'n_tokens_neox': 1607, 'n_words': 623}