Coverage for src/loman/visualization.py: 100%
408 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-09-07 00:32 +0000
1"""Visualization tools for computation graphs using Graphviz."""
3import os
4import subprocess # nosec B404
5import sys
6import tempfile
7from abc import ABC, abstractmethod
8from collections import defaultdict
9from collections.abc import Sequence
10from dataclasses import dataclass, field
11from typing import TYPE_CHECKING, Any, ClassVar, Protocol, runtime_checkable
13import matplotlib as mpl
14import networkx as nx
15import numpy as np
16import pandas as pd
17import pydotplus
18from matplotlib.colors import Colormap
20from .consts import NodeAttributes, NodeTransformations, States
21from .graph_utils import contract_node
22from .nodekey import Name, NodeKey, is_pattern, match_pattern, to_nodekey
24if TYPE_CHECKING:
25 from .computeengine import Computation
28@runtime_checkable
29class _ComputationLike(Protocol):
30 """Structural stand-in for :class:`~loman.computeengine.Computation`.
32 ``visualization`` is a lower layer than ``computeengine`` (which imports
33 :class:`GraphView`/:class:`NodeFormatter` from here), so importing the
34 top-level ``loman`` package — or ``computeengine`` — at module load would
35 create an import cycle. A nested computation stored as a node value is
36 instead detected structurally, by the attributes that identify a
37 Computation, without importing anything from the upper layer.
38 """
40 dag: Any
42 def add_node(self, *args: Any, **kwargs: Any) -> Any:
43 """Marker method present on any Computation."""
44 ...
46 def compute(self, *args: Any, **kwargs: Any) -> Any:
47 """Marker method present on any Computation."""
48 ...
51@dataclass
52class Node:
53 """Represents a node in the visualization graph."""
55 nodekey: NodeKey
56 original_nodekey: NodeKey
57 data: dict[str, Any]
60class NodeFormatter(ABC):
61 """Abstract base class for node formatting in visualizations."""
63 @abstractmethod
64 def calibrate(self, nodes: list[Node]) -> None:
65 """Calibrate formatter based on all nodes in the graph."""
66 # pragma: no cover
68 @abstractmethod
69 def format(self, name: NodeKey, nodes: list[Node], is_composite: bool) -> dict[str, Any] | None:
70 """Format node appearance returning dict of graphviz attributes."""
71 # pragma: no cover
73 @staticmethod
74 def create(
75 cmap: dict[States | None, str] | Colormap | None = None, colors: str = "state", shapes: str | None = None
76 ) -> "CompositeNodeFormatter":
77 """Create a composite node formatter with specified color and shape options."""
78 node_formatters: list[NodeFormatter] = [StandardLabel(), StandardGroup()]
80 if isinstance(shapes, str):
81 shapes = shapes.lower()
82 if shapes == "type":
83 node_formatters.append(ShapeByType())
84 elif shapes is None:
85 pass
86 else:
87 msg = f"{shapes} is not a valid loman shapes parameter for visualization"
88 raise ValueError(msg)
90 colors = colors.lower()
91 if colors == "state":
92 state_cmap = cmap if isinstance(cmap, dict) else None
93 node_formatters.append(ColorByState(state_cmap)) # type: ignore[arg-type]
94 elif colors == "timing":
95 timing_cmap = cmap if isinstance(cmap, Colormap) else None
96 node_formatters.append(ColorByTiming(timing_cmap))
97 else:
98 msg = f"{colors} is not a valid loman colors parameter for visualization"
99 raise ValueError(msg)
101 node_formatters.append(StandardStylingOverrides())
102 node_formatters.append(RectBlocks())
104 return CompositeNodeFormatter(node_formatters)
107def aggregate_states(states: Sequence[States | None]) -> States | None:
108 """Reduce the states of one rendered node's members to a single state.
110 A rendered node stands for either one computation node or, when a block is
111 collapsed, all of its members. A single member reports its own state.
112 Otherwise ERROR wins, then STALE, then the common state if every member
113 agrees, and a genuine mixture reduces to ``None``.
115 ``None`` is Loman's mixed marker throughout: :class:`ColorByState` paints it
116 white and the notebook widget labels it ``MIXED``. Both read it from here so
117 that the rendered colour and the reported state cannot drift apart.
119 :param states: States of the members behind one rendered node.
120 :return: The state to display, or ``None`` when the members disagree.
121 """
122 if len(states) == 1:
123 return states[0]
124 if any(s == States.ERROR for s in states):
125 return States.ERROR
126 if any(s == States.STALE for s in states):
127 return States.STALE
128 first = states[0] if states else None
129 return first if all(s == first for s in states) else None
132class ColorByState(NodeFormatter):
133 """Node formatter that colors nodes based on their computation state."""
135 DEFAULT_STATE_COLORS: ClassVar[dict[States | None, str]] = {
136 None: "#ffffff", # xkcd white
137 States.PLACEHOLDER: "#f97306", # xkcd orange
138 States.UNINITIALIZED: "#0343df", # xkcd blue
139 States.STALE: "#ffff14", # xkcd yellow
140 States.COMPUTABLE: "#9dff00", # xkcd bright yellow green
141 States.UPTODATE: "#15b01a", # xkcd green
142 States.ERROR: "#e50000", # xkcd red
143 States.PINNED: "#bf77f6", # xkcd light purple
144 }
146 def __init__(self, state_colors: dict[States | None, str] | None = None) -> None:
147 """Initialize with custom state color mapping."""
148 if state_colors is None:
149 state_colors = self.DEFAULT_STATE_COLORS.copy()
150 self.state_colors = state_colors
152 def calibrate(self, nodes: list[Node]) -> None:
153 """Calibrate formatter based on all nodes in the graph."""
155 def format(self, name: NodeKey, nodes: list[Node], is_composite: bool) -> dict[str, Any] | None:
156 """Format node color based on computation state."""
157 state = aggregate_states([node.data.get(NodeAttributes.STATE, None) for node in nodes])
158 return {"style": "filled", "fillcolor": self.state_colors[state]}
161class ColorByTiming(NodeFormatter):
162 """Node formatter that colors nodes based on their execution timing."""
164 def __init__(self, cmap: Colormap | None = None) -> None:
165 """Initialize with an optional colormap for timing visualization."""
166 if cmap is None:
167 cmap = mpl.colors.LinearSegmentedColormap.from_list("blend", ["#15b01a", "#ffff14", "#e50000"])
168 self.cmap = cmap
169 self.min_duration: float = float("nan")
170 self.max_duration: float = float("nan")
172 def calibrate(self, nodes: list[Node]) -> None:
173 """Calibrate the color mapping based on node timing data."""
174 durations: list[float] = []
175 for node in nodes:
176 timing = node.data.get(NodeAttributes.TIMING)
177 if timing is not None:
178 durations.append(timing.duration)
179 if durations:
180 self.max_duration = max(durations)
181 self.min_duration = min(durations)
183 def format(self, name: NodeKey, nodes: list[Node], is_composite: bool) -> dict[str, Any] | None:
184 """Format a node with timing-based coloring."""
185 if len(nodes) == 1:
186 data = nodes[0].data
187 timing_data = data.get(NodeAttributes.TIMING)
188 if timing_data is None:
189 col = "#FFFFFF"
190 else:
191 duration = timing_data.duration
192 norm_duration: float = (duration - self.min_duration) / max(1e-8, self.max_duration - self.min_duration)
193 col = mpl.colors.rgb2hex(self.cmap(norm_duration))
194 return {"style": "filled", "fillcolor": col}
195 return None
198class ShapeByType(NodeFormatter):
199 """Node formatter that sets node shapes based on their type."""
201 def calibrate(self, nodes: list[Node]) -> None:
202 """Calibrate formatter based on all nodes in the graph."""
204 def format(self, name: NodeKey, nodes: list[Node], is_composite: bool) -> dict[str, Any] | None:
205 """Format a node with type-based shape styling."""
206 if len(nodes) == 1:
207 data = nodes[0].data
208 value = data.get(NodeAttributes.VALUE)
209 if value is None:
210 return None
211 if isinstance(value, np.ndarray):
212 return {"shape": "rect"}
213 elif isinstance(value, pd.DataFrame):
214 return {"shape": "box3d"}
215 elif np.isscalar(value):
216 return {"shape": "ellipse"}
217 elif isinstance(value, (list, tuple)):
218 return {"shape": "ellipse", "peripheries": 2}
219 elif isinstance(value, dict):
220 return {"shape": "house", "peripheries": 2}
221 elif isinstance(value, _ComputationLike):
222 return {"shape": "hexagon"}
223 else:
224 return {"shape": "diamond"}
225 return None
228class RectBlocks(NodeFormatter):
229 """Node formatter that shapes composite nodes as rectangles."""
231 def calibrate(self, nodes: list[Node]) -> None:
232 """Calibrate formatter based on all nodes in the graph."""
234 def format(self, name: NodeKey, nodes: list[Node], is_composite: bool) -> dict[str, Any] | None:
235 """Return rectangle shape for composite nodes."""
236 if is_composite:
237 return {"shape": "rect", "peripheries": 2}
238 return None
241class StandardLabel(NodeFormatter):
242 """Node formatter that sets node labels."""
244 def calibrate(self, nodes: list[Node]) -> None:
245 """Calibrate formatter based on all nodes in the graph."""
247 def format(self, name: NodeKey, nodes: list[Node], is_composite: bool) -> dict[str, Any] | None:
248 """Return standard label for node."""
249 return {"label": name.label}
252def get_group_path(name: NodeKey, data: dict[str, Any]) -> NodeKey:
253 """Determine the group path for a node based on name hierarchy and group attribute."""
254 name_group_path = name.parent
255 attribute_group = data.get(NodeAttributes.GROUP)
256 attribute_group_path = None if attribute_group is None else NodeKey((attribute_group,))
258 group_path = name_group_path.join(attribute_group_path)
259 return group_path
262class StandardGroup(NodeFormatter):
263 """Node formatter that applies standard grouping styles."""
265 def calibrate(self, nodes: list[Node]) -> None:
266 """Calibrate formatter based on all nodes in the graph."""
268 def format(self, name: NodeKey, nodes: list[Node], is_composite: bool) -> dict[str, Any] | None:
269 """Format a node with standard group styling."""
270 if len(nodes) == 1:
271 data = nodes[0].data
272 group_path = get_group_path(name, data)
273 else:
274 group_path = name.parent
275 if group_path.is_root:
276 return None
277 return {"_group": group_path}
280class StandardStylingOverrides(NodeFormatter):
281 """Node formatter that applies standard styling overrides."""
283 def calibrate(self, nodes: list[Node]) -> None:
284 """Calibrate formatter based on all nodes in the graph."""
286 def format(self, name: NodeKey, nodes: list[Node], is_composite: bool) -> dict[str, Any] | None:
287 """Format a node with standard styling overrides."""
288 if len(nodes) == 1:
289 data = nodes[0].data
290 style = data.get(NodeAttributes.STYLE)
291 if style is None:
292 return None
293 if style == "small":
294 return {"width": 0.3, "height": 0.2, "fontsize": 8}
295 elif style == "dot":
296 return {"shape": "point", "width": 0.1, "peripheries": 1}
297 return None
300@dataclass
301class CompositeNodeFormatter(NodeFormatter):
302 """A node formatter that combines multiple formatters together."""
304 formatters: list[NodeFormatter] = field(default_factory=list)
306 def calibrate(self, nodes: list[Node]) -> None:
307 """Calibrate all the contained formatters with the given nodes."""
308 for formatter in self.formatters:
309 formatter.calibrate(nodes)
311 def format(self, name: NodeKey, nodes: list[Node], is_composite: bool) -> dict[str, Any] | None:
312 """Format a node by combining output from all contained formatters."""
313 d: dict[str, Any] = {}
314 for formatter in self.formatters:
315 format_attrs = formatter.format(name, nodes, is_composite)
316 if format_attrs is not None:
317 d.update(format_attrs)
318 return d
321@dataclass
322class GraphView:
323 """A view for visualizing computation graphs as graphical diagrams."""
325 computation: "Computation"
326 root: Name | None = None
327 node_formatter: NodeFormatter | None = None
328 node_transformations: dict[Name, str] | None = None
329 collapse_all: bool = True
331 graph_attr: dict[str, Any] | None = None
332 node_attr: dict[str, Any] | None = None
333 edge_attr: dict[str, Any] | None = None
335 struct_dag: nx.DiGraph | None = None
336 viz_dag: nx.DiGraph | None = None
337 viz_dot: pydotplus.Dot | None = None
338 original_nodes: defaultdict[NodeKey, list[NodeKey]] = field(default_factory=lambda: defaultdict(list), init=False)
339 composite_nodes: set[NodeKey] = field(default_factory=set, init=False)
340 node_index_map: dict[NodeKey, str] = field(default_factory=dict, init=False)
342 def __post_init__(self) -> None:
343 """Initialize the graph view after dataclass construction."""
344 self.refresh()
346 @staticmethod
347 def get_sub_block(
348 dag: nx.DiGraph, root: Name | None, node_transformations: dict[NodeKey, str]
349 ) -> tuple[nx.DiGraph, defaultdict[NodeKey, list[NodeKey]], set[NodeKey]]:
350 """Extract a subgraph with node transformations for visualization."""
351 d_transform_to_nodes: defaultdict[str, list[NodeKey]] = defaultdict(list)
352 for nk, transform in node_transformations.items():
353 d_transform_to_nodes[transform].append(nk)
355 dag_out: nx.DiGraph = nx.DiGraph()
357 d_original_to_mapped: dict[NodeKey, NodeKey] = {}
358 s_collapsed: set[NodeKey] = set()
360 for nk_original in dag.nodes():
361 nk_mapped = nk_original.drop_root(root)
362 if nk_mapped is None:
363 continue
364 nk_highest_collapse = nk_original
365 is_collapsed = False
366 for nk_collapse in d_transform_to_nodes[NodeTransformations.COLLAPSE]:
367 if nk_highest_collapse.is_descendent_of(nk_collapse):
368 nk_highest_collapse = nk_collapse
369 is_collapsed = True
370 nk_mapped = nk_highest_collapse.drop_root(root)
371 if nk_mapped is None: # pragma: no cover
372 continue
373 d_original_to_mapped[nk_original] = nk_mapped
374 if is_collapsed:
375 s_collapsed.add(nk_mapped)
377 for nk_mapped in d_original_to_mapped.values():
378 dag_out.add_node(nk_mapped)
380 for nk_u, nk_v in dag.edges():
381 nk_mapped_u = d_original_to_mapped.get(nk_u)
382 nk_mapped_v = d_original_to_mapped.get(nk_v)
383 if nk_mapped_u is None or nk_mapped_v is None or nk_mapped_u == nk_mapped_v:
384 continue
385 dag_out.add_edge(nk_mapped_u, nk_mapped_v)
387 for nk in d_transform_to_nodes[NodeTransformations.CONTRACT]:
388 contract_node(dag_out, d_original_to_mapped[nk])
389 del d_original_to_mapped[nk]
391 d_mapped_to_original: defaultdict[NodeKey, list[NodeKey]] = defaultdict(list)
392 for nk_original, nk_mapped in d_original_to_mapped.items():
393 if nk_mapped in dag_out.nodes:
394 d_mapped_to_original[nk_mapped].append(nk_original)
396 s_collapsed.intersection_update(dag_out.nodes)
398 return dag_out, d_mapped_to_original, s_collapsed
400 def _initialize_transforms(self) -> dict[NodeKey, str]:
401 """Initialize node transformations for visualization."""
402 node_transformations: dict[NodeKey, str] = {}
403 if self.collapse_all:
404 self._apply_default_collapse_transforms(node_transformations)
405 self._apply_custom_transforms(node_transformations)
406 return node_transformations
408 def _apply_default_collapse_transforms(self, node_transformations: dict[NodeKey, str]) -> None:
409 """Apply default collapse transformations to tree nodes."""
410 for n in self.computation.get_tree_descendents(self.root):
411 nk = to_nodekey(n)
412 if not self.computation.has_node(nk):
413 node_transformations[nk] = NodeTransformations.COLLAPSE
415 def _apply_custom_transforms(self, node_transformations: dict[NodeKey, str]) -> dict[NodeKey, str]:
416 """Apply user-specified custom transformations to nodes."""
417 if self.node_transformations is not None:
418 for rule_name, transform in self.node_transformations.items():
419 include_ancestors = transform == NodeTransformations.EXPAND
420 rule_nk = to_nodekey(rule_name)
421 if is_pattern(rule_nk):
422 apply_nodes: set[NodeKey] = set()
423 for n in self.computation.get_tree_descendents(self.root):
424 nk = to_nodekey(n)
425 if match_pattern(rule_nk, nk):
426 apply_nodes.add(nk)
427 else:
428 apply_nodes = {rule_nk}
429 node_transformations[rule_nk] = transform
430 if include_ancestors:
431 for nk in apply_nodes:
432 for nk1 in nk.ancestors():
433 if nk1.is_root or nk1 == self.root:
434 break
435 node_transformations[nk1] = NodeTransformations.EXPAND
436 for r_nk in apply_nodes:
437 node_transformations[r_nk] = transform
438 return node_transformations
440 def _create_visualization_dag(
441 self,
442 original_nodes: defaultdict[NodeKey, list[NodeKey]],
443 composite_nodes: set[NodeKey],
444 node_index_map: dict[NodeKey, str],
445 ) -> nx.DiGraph:
446 """Create the visualization DAG from structure and node data."""
447 node_formatter = self.node_formatter
448 if node_formatter is None:
449 node_formatter = NodeFormatter.create()
450 assert self.struct_dag is not None # noqa: S101
451 return create_viz_dag(
452 self.struct_dag,
453 self.computation.dag,
454 node_formatter,
455 original_nodes,
456 composite_nodes,
457 node_index_map=node_index_map,
458 )
460 def _create_dot_graph(self) -> pydotplus.Dot:
461 """Create a PyDot graph from the visualization DAG."""
462 return to_pydot(self.viz_dag, self.graph_attr, self.node_attr, self.edge_attr)
464 def refresh(self) -> None:
465 """Refresh the visualization by rebuilding the graph structure."""
466 node_transformations = self._initialize_transforms()
467 self.struct_dag, self.original_nodes, self.composite_nodes = self.get_sub_block(
468 self.computation.dag, self.root, node_transformations
469 )
470 self.node_index_map = {}
471 self.viz_dag = self._create_visualization_dag(self.original_nodes, self.composite_nodes, self.node_index_map)
472 self.viz_dot = self._create_dot_graph()
474 def svg(self) -> str | None:
475 """Generate SVG representation of the visualization."""
476 if self.viz_dot is None:
477 return None
478 svg_bytes: bytes = self.viz_dot.create_svg() # type: ignore[attr-defined]
479 return svg_bytes.decode("utf-8")
481 def view(self) -> None: # pragma: no cover
482 """Open the visualization in a PDF viewer."""
483 assert self.viz_dot is not None # noqa: S101
484 with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as f:
485 f.write(self.viz_dot.create_pdf()) # type: ignore[attr-defined]
486 if sys.platform == "win32":
487 os.startfile(f.name) # pragma: no cover # nosec B606 # noqa: S606
488 else:
489 subprocess.run(["open", f.name], check=False) # pragma: no cover # nosec B603 B607 # noqa: S603, S607
491 def _repr_svg_(self) -> str | None:
492 """Return SVG representation for Jupyter notebook display."""
493 return self.svg()
496def create_viz_dag(
497 struct_dag: nx.DiGraph,
498 comp_dag: nx.DiGraph,
499 node_formatter: NodeFormatter,
500 original_nodes: defaultdict[NodeKey, list[NodeKey]],
501 composite_nodes: set[NodeKey],
502 node_index_map: dict[NodeKey, str] | None = None,
503) -> nx.DiGraph:
504 """Create a visualization DAG from the computation structure.
506 :param struct_dag: Structural graph of the nodes actually being rendered.
507 :param comp_dag: The full computation graph, for node data.
508 :param node_formatter: Formatter supplying each node's Graphviz attributes.
509 :param original_nodes: Members standing behind each rendered node.
510 :param composite_nodes: Rendered nodes that stand for a collapsed block.
511 :param node_index_map: Optional dict to fill in with the node-key to
512 rendered-name mapping. The mapping is assigned by enumeration order and
513 cannot be reconstructed afterwards, so callers that need node identity
514 --- the notebook widget, via :attr:`GraphView.node_index_map` --- pass a
515 dict in rather than duplicating the numbering rule.
516 :return: The visualization DAG, ready to convert to Graphviz.
517 """
518 if node_formatter is not None:
519 nodes: list[Node] = []
520 for nodekey in struct_dag.nodes:
521 for original_nodekey in original_nodes[nodekey]:
522 data = comp_dag.nodes[original_nodekey]
523 n = Node(nodekey, original_nodekey, data)
524 nodes.append(n)
525 node_formatter.calibrate(nodes)
527 viz_dag: nx.DiGraph = nx.DiGraph()
528 if node_index_map is None:
529 node_index_map = {}
530 for i, nodekey in enumerate(struct_dag.nodes):
531 short_name = f"n{i}"
532 attr_dict: dict[str, Any] | None = None
534 if node_formatter is not None:
535 nodes = []
536 for original_nodekey in original_nodes[nodekey]:
537 data = comp_dag.nodes[original_nodekey]
538 n = Node(nodekey, original_nodekey, data)
539 nodes.append(n)
540 is_composite = nodekey in composite_nodes
541 attr_dict = node_formatter.format(nodekey, nodes, is_composite)
542 if attr_dict is None: # pragma: no cover
543 attr_dict = {}
545 attr_dict = {k: v for k, v in attr_dict.items() if v is not None}
547 viz_dag.add_node(short_name, **attr_dict)
548 node_index_map[nodekey] = short_name
550 for name1, name2 in struct_dag.edges():
551 short_name_1 = node_index_map[name1]
552 short_name_2 = node_index_map[name2]
554 group_path1 = get_group_path(name1, struct_dag.nodes[name1])
555 group_path2 = get_group_path(name2, struct_dag.nodes[name2])
556 group_path = NodeKey.common_parent(group_path1, group_path2)
558 edge_attr_dict: dict[str, Any] = {}
559 if not group_path.is_root:
560 # group_path = None
561 edge_attr_dict["_group"] = group_path
563 viz_dag.add_edge(short_name_1, short_name_2, **edge_attr_dict)
565 return viz_dag
568def _group_nodes_and_edges(
569 viz_dag: nx.DiGraph,
570) -> tuple[NodeKey, dict[NodeKey, list[str]], dict[NodeKey, list[tuple[str, str]]]]:
571 """Group nodes and edges by their groups."""
572 root = NodeKey.root()
574 node_groups: dict[NodeKey, list[str]] = {}
575 for name, data in viz_dag.nodes(data=True):
576 group = data.get("_group", root)
577 node_groups.setdefault(group, []).append(name)
579 edge_groups: dict[NodeKey, list[tuple[str, str]]] = {}
580 for name1, name2, data in viz_dag.edges(data=True):
581 group = data.get("_group", root)
582 edge_groups.setdefault(group, []).append((name1, name2))
584 return root, node_groups, edge_groups
587def _create_pydot_nodes(
588 viz_dag: nx.DiGraph,
589 node_groups: dict[NodeKey, list[str]],
590 subgraphs: dict[NodeKey, pydotplus.Dot | pydotplus.Subgraph],
591 root: NodeKey,
592) -> None:
593 """Create PyDot nodes for each group."""
594 for group, names in node_groups.items():
595 c = subgraphs[root] if group is root else create_subgraph(group)
597 for name in names:
598 node = pydotplus.Node(name)
599 for k, v in viz_dag.nodes[name].items():
600 if not k.startswith("_"):
601 node.set(k, v)
602 c.add_node(node)
604 subgraphs[group] = c
607def _ensure_parent_subgraphs(subgraphs: dict[NodeKey, pydotplus.Dot | pydotplus.Subgraph]) -> None:
608 """Ensure all parent subgraphs exist in the hierarchy."""
609 groups = list(subgraphs.keys())
610 for group in groups:
611 group1 = group
612 while True:
613 if group1.is_root:
614 break
615 group1 = group1.parent
616 if group1 in subgraphs:
617 break
618 subgraphs[group1] = create_subgraph(group1)
621def _link_subgraphs(subgraphs: dict[NodeKey, pydotplus.Dot | pydotplus.Subgraph]) -> None:
622 """Link subgraphs to their parents."""
623 for group, subgraph in subgraphs.items():
624 if group.is_root:
625 continue
626 parent = group
627 while True:
628 parent = parent.parent
629 if parent in subgraphs or parent.is_root:
630 break
631 subgraphs[parent].add_subgraph(subgraph)
634def _add_edges_to_subgraphs(
635 edge_groups: dict[NodeKey, list[tuple[str, str]]], subgraphs: dict[NodeKey, pydotplus.Dot | pydotplus.Subgraph]
636) -> None:
637 """Add edges to their respective subgraphs."""
638 for group, edges in edge_groups.items():
639 c = subgraphs[group]
640 for name1, name2 in edges:
641 edge = pydotplus.Edge(name1, name2)
642 c.add_edge(edge)
645def to_pydot(
646 viz_dag: nx.DiGraph | None,
647 graph_attr: dict[str, Any] | None = None,
648 node_attr: dict[str, Any] | None = None,
649 edge_attr: dict[str, Any] | None = None,
650) -> pydotplus.Dot:
651 """Convert a visualization DAG to a PyDot graph for rendering."""
652 assert viz_dag is not None # noqa: S101
653 root, node_groups, edge_groups = _group_nodes_and_edges(viz_dag)
655 subgraphs: dict[NodeKey, pydotplus.Dot | pydotplus.Subgraph] = {
656 root: create_root_graph(graph_attr, node_attr, edge_attr)
657 }
659 _create_pydot_nodes(viz_dag, node_groups, subgraphs, root)
660 _ensure_parent_subgraphs(subgraphs)
661 _link_subgraphs(subgraphs)
662 _add_edges_to_subgraphs(edge_groups, subgraphs)
664 result = subgraphs[root]
665 assert isinstance(result, pydotplus.Dot) # noqa: S101
666 return result
669def create_root_graph(
670 graph_attr: dict[str, Any] | None, node_attr: dict[str, Any] | None, edge_attr: dict[str, Any] | None
671) -> pydotplus.Dot:
672 """Create root Graphviz graph with specified attributes.
674 Notes:
675 Graphviz attributes like size expect a quoted string when containing
676 commas (e.g. "10,8"). Some pydotplus setters don't auto-quote, which
677 can produce a DOT syntax error near ',' if we pass a raw string.
678 We defensively quote string values that contain commas or whitespace.
679 """
681 def _normalize_attr_value(v: Any) -> Any:
682 """Normalize attribute values for Graphviz, quoting strings as needed."""
683 # Keep numeric values as-is
684 if isinstance(v, (int, float)):
685 return v
686 s = str(v)
687 # If already quoted, keep
688 if len(s) >= 2 and ((s[0] == '"' and s[-1] == '"') or (s[0] == "'" and s[-1] == "'")):
689 return s
690 # Quote if contains comma, whitespace, or special characters
691 if any(c in s for c in [",", " ", "\t", "\n"]) or s == "":
692 return f'"{s}"'
693 return s
695 root_graph = pydotplus.Dot()
696 if graph_attr is not None:
697 for k, v in graph_attr.items():
698 root_graph.set(k, _normalize_attr_value(v))
699 if node_attr is not None:
700 # For node/edge defaults, normalize each value too
701 node_defaults = {k: _normalize_attr_value(v) for k, v in node_attr.items()}
702 root_graph.set_node_defaults(**node_defaults)
703 if edge_attr is not None:
704 edge_defaults = {k: _normalize_attr_value(v) for k, v in edge_attr.items()}
705 root_graph.set_edge_defaults(**edge_defaults)
706 return root_graph
709def create_subgraph(group: NodeKey) -> pydotplus.Subgraph:
710 """Create a Graphviz subgraph for a node group."""
711 c = pydotplus.Subgraph("cluster_" + str(group))
712 c.obj_dict["attributes"]["label"] = str(group)
713 return c