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