1import os
2import re
3from functools import partial
4
5from .core import get_dependencies, ishashable, istask
6from .utils import apply, funcname, import_required, key_split
7
8graphviz = import_required(
9    "graphviz",
10    "Drawing dask graphs requires the "
11    "`graphviz` python library and the "
12    "`graphviz` system library to be "
13    "installed.",
14)
15
16
17def task_label(task):
18    """Label for a task on a dot graph.
19
20    Examples
21    --------
22    >>> from operator import add
23    >>> task_label((add, 1, 2))
24    'add'
25    >>> task_label((add, (add, 1, 2), 3))
26    'add(...)'
27    """
28    func = task[0]
29    if func is apply:
30        func = task[1]
31    if hasattr(func, "funcs"):
32        if len(func.funcs) > 1:
33            return f"{funcname(func.funcs[0])}(...)"
34        else:
35            head = funcname(func.funcs[0])
36    else:
37        head = funcname(func)
38    if any(has_sub_tasks(i) for i in task[1:]):
39        return f"{head}(...)"
40    else:
41        return head
42
43
44def has_sub_tasks(task):
45    """Returns True if the task has sub tasks"""
46    if istask(task):
47        return True
48    elif isinstance(task, list):
49        return any(has_sub_tasks(i) for i in task)
50    else:
51        return False
52
53
54def name(x):
55    try:
56        return str(hash(x))
57    except TypeError:
58        return str(hash(str(x)))
59
60
61_HASHPAT = re.compile("([0-9a-z]{32})")
62_UUIDPAT = re.compile("([0-9a-z]{8}-[0-9a-z]{4}-[0-9a-z]{4}-[0-9a-z]{4}-[0-9a-z]{12})")
63
64
65def label(x, cache=None):
66    """
67
68    >>> label('x')
69    'x'
70
71    >>> label(('x', 1))
72    "('x', 1)"
73
74    >>> from hashlib import md5
75    >>> x = 'x-%s-hello' % md5(b'1234').hexdigest()
76    >>> x
77    'x-81dc9bdb52d04dc20036dbd8313ed055-hello'
78
79    >>> label(x)
80    'x-#-hello'
81
82    >>> from uuid import uuid1
83    >>> x = 'x-%s-hello' % uuid1()
84    >>> x  # doctest: +SKIP
85    'x-4c1a3d7e-0b45-11e6-8334-54ee75105593-hello'
86
87    >>> label(x)
88    'x-#-hello'
89    """
90    s = str(x)
91    for pattern in (_HASHPAT, _UUIDPAT):
92        m = re.search(pattern, s)
93        if m is not None:
94            for h in m.groups():
95                if cache is not None:
96                    n = cache.get(h, len(cache))
97                    label = f"#{n}"
98                    # cache will be overwritten destructively
99                    cache[h] = n
100                else:
101                    label = "#"
102                s = s.replace(h, label)
103    return s
104
105
106def box_label(key, verbose=False):
107    """Label boxes in graph by chunk index
108
109    >>> box_label(('x', 1, 2, 3))
110    '(1, 2, 3)'
111    >>> box_label(('x', 123))
112    '123'
113    >>> box_label('x')
114    ''
115    """
116    if isinstance(key, tuple):
117        key = key[1:]
118        if len(key) == 1:
119            [key] = key
120        return str(key)
121    elif verbose:
122        return str(key)
123    else:
124        return ""
125
126
127def to_graphviz(
128    dsk,
129    data_attributes=None,
130    function_attributes=None,
131    rankdir="BT",
132    graph_attr=None,
133    node_attr=None,
134    edge_attr=None,
135    collapse_outputs=False,
136    verbose=False,
137    **kwargs,
138):
139    data_attributes = data_attributes or {}
140    function_attributes = function_attributes or {}
141    graph_attr = graph_attr or {}
142    node_attr = node_attr or {}
143    edge_attr = edge_attr or {}
144
145    graph_attr["rankdir"] = rankdir
146    node_attr["fontname"] = "helvetica"
147
148    graph_attr.update(kwargs)
149    g = graphviz.Digraph(
150        graph_attr=graph_attr, node_attr=node_attr, edge_attr=edge_attr
151    )
152
153    seen = set()
154    connected = set()
155
156    for k, v in dsk.items():
157        k_name = name(k)
158        if istask(v):
159            func_name = name((k, "function")) if not collapse_outputs else k_name
160            if collapse_outputs or func_name not in seen:
161                seen.add(func_name)
162                attrs = function_attributes.get(k, {}).copy()
163                attrs.setdefault("label", key_split(k))
164                attrs.setdefault("shape", "circle")
165                g.node(func_name, **attrs)
166            if not collapse_outputs:
167                g.edge(func_name, k_name)
168                connected.add(func_name)
169                connected.add(k_name)
170
171            for dep in get_dependencies(dsk, k):
172                dep_name = name(dep)
173                if dep_name not in seen:
174                    seen.add(dep_name)
175                    attrs = data_attributes.get(dep, {}).copy()
176                    attrs.setdefault("label", box_label(dep, verbose))
177                    attrs.setdefault("shape", "box")
178                    g.node(dep_name, **attrs)
179                g.edge(dep_name, func_name)
180                connected.add(dep_name)
181                connected.add(func_name)
182
183        elif ishashable(v) and v in dsk:
184            v_name = name(v)
185            g.edge(v_name, k_name)
186            connected.add(v_name)
187            connected.add(k_name)
188
189        if (not collapse_outputs or k_name in connected) and k_name not in seen:
190            seen.add(k_name)
191            attrs = data_attributes.get(k, {}).copy()
192            attrs.setdefault("label", box_label(k, verbose))
193            attrs.setdefault("shape", "box")
194            g.node(k_name, **attrs)
195    return g
196
197
198IPYTHON_IMAGE_FORMATS = frozenset(["jpeg", "png"])
199IPYTHON_NO_DISPLAY_FORMATS = frozenset(["dot", "pdf"])
200
201
202def _get_display_cls(format):
203    """
204    Get the appropriate IPython display class for `format`.
205
206    Returns `IPython.display.SVG` if format=='svg', otherwise
207    `IPython.display.Image`.
208
209    If IPython is not importable, return dummy function that swallows its
210    arguments and returns None.
211    """
212    dummy = lambda *args, **kwargs: None
213    try:
214        import IPython.display as display
215    except ImportError:
216        # Can't return a display object if no IPython.
217        return dummy
218
219    if format in IPYTHON_NO_DISPLAY_FORMATS:
220        # IPython can't display this format natively, so just return None.
221        return dummy
222    elif format in IPYTHON_IMAGE_FORMATS:
223        # Partially apply `format` so that `Image` and `SVG` supply a uniform
224        # interface to the caller.
225        return partial(display.Image, format=format)
226    elif format == "svg":
227        return display.SVG
228    else:
229        raise ValueError("Unknown format '%s' passed to `dot_graph`" % format)
230
231
232def dot_graph(dsk, filename="mydask", format=None, **kwargs):
233    """
234    Render a task graph using dot.
235
236    If `filename` is not None, write a file to disk with the specified name and extension.
237    If no extension is specified, '.png' will be used by default.
238
239    Parameters
240    ----------
241    dsk : dict
242        The graph to display.
243    filename : str or None, optional
244        The name of the file to write to disk. If the provided `filename`
245        doesn't include an extension, '.png' will be used by default.
246        If `filename` is None, no file will be written, and we communicate
247        with dot using only pipes.  Default is 'mydask'.
248    format : {'png', 'pdf', 'dot', 'svg', 'jpeg', 'jpg'}, optional
249        Format in which to write output file.  Default is 'png'.
250    **kwargs
251        Additional keyword arguments to forward to `to_graphviz`.
252
253    Returns
254    -------
255    result : None or IPython.display.Image or IPython.display.SVG  (See below.)
256
257    Notes
258    -----
259    If IPython is installed, we return an IPython.display object in the
260    requested format.  If IPython is not installed, we just return None.
261
262    We always return None if format is 'pdf' or 'dot', because IPython can't
263    display these formats natively. Passing these formats with filename=None
264    will not produce any useful output.
265
266    See Also
267    --------
268    dask.dot.to_graphviz
269    """
270    g = to_graphviz(dsk, **kwargs)
271    return graphviz_to_file(g, filename, format)
272
273
274def graphviz_to_file(g, filename, format):
275    fmts = [".png", ".pdf", ".dot", ".svg", ".jpeg", ".jpg"]
276
277    if (
278        format is None
279        and filename is not None
280        and any(filename.lower().endswith(fmt) for fmt in fmts)
281    ):
282        filename, format = os.path.splitext(filename)
283        format = format[1:].lower()
284
285    if format is None:
286        format = "png"
287
288    data = g.pipe(format=format)
289    if not data:
290        raise RuntimeError(
291            "Graphviz failed to properly produce an image. "
292            "This probably means your installation of graphviz "
293            "is missing png support. See: "
294            "https://github.com/ContinuumIO/anaconda-issues/"
295            "issues/485 for more information."
296        )
297
298    display_cls = _get_display_cls(format)
299
300    if filename is None:
301        return display_cls(data=data)
302
303    full_filename = ".".join([filename, format])
304    with open(full_filename, "wb") as f:
305        f.write(data)
306
307    return display_cls(filename=full_filename)
308