Coverage for src/loman/visualization.py: 100%
407 statements
« prev ^ index » next coverage.py v7.15.1, created at 2026-07-13 06:52 +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 dataclasses import dataclass, field
10from typing import TYPE_CHECKING, Any, ClassVar, Protocol, runtime_checkable
12import matplotlib as mpl
13import networkx as nx
14import numpy as np
15import pandas as pd
16import pydotplus
17from matplotlib.colors import Colormap
19from .consts import NodeAttributes, NodeTransformations, States
20from .graph_utils import contract_node
21from .nodekey import Name, NodeKey, is_pattern, match_pattern, to_nodekey
23if TYPE_CHECKING:
24 from .computeengine import Computation
27@runtime_checkable
28class _ComputationLike(Protocol):
29 """Structural stand-in for :class:`~loman.computeengine.Computation`.
31 ``visualization`` is a lower layer than ``computeengine`` (which imports
32 :class:`GraphView`/:class:`NodeFormatter` from here), so importing the
33 top-level ``loman`` package — or ``computeengine`` — at module load would
34 create an import cycle. A nested computation stored as a node value is
35 instead detected structurally, by the attributes that identify a
36 Computation, without importing anything from the upper layer.
37 """
39 dag: Any
41 def add_node(self, *args: Any, **kwargs: Any) -> Any:
42 """Marker method present on any Computation."""
43 ...
45 def compute(self, *args: Any, **kwargs: Any) -> Any:
46 """Marker method present on any Computation."""
47 ...
50@dataclass
51class Node:
52 """Represents a node in the visualization graph."""
54 nodekey: NodeKey
55 original_nodekey: NodeKey
56 data: dict[str, Any]
59class NodeFormatter(ABC):
60 """Abstract base class for node formatting in visualizations."""
62 @abstractmethod
63 def calibrate(self, nodes: list[Node]) -> None:
64 """Calibrate formatter based on all nodes in the graph."""
65 pass # pragma: no cover
67 @abstractmethod
68 def format(self, name: NodeKey, nodes: list[Node], is_composite: bool) -> dict[str, Any] | None:
69 """Format node appearance returning dict of graphviz attributes."""
70 pass # pragma: no cover
72 @staticmethod
73 def create(
74 cmap: dict[States | None, str] | Colormap | None = None, colors: str = "state", shapes: str | None = None
75 ) -> "CompositeNodeFormatter":
76 """Create a composite node formatter with specified color and shape options."""
77 node_formatters: list[NodeFormatter] = [StandardLabel(), StandardGroup()]
79 if isinstance(shapes, str):
80 shapes = shapes.lower()
81 if shapes == "type":
82 node_formatters.append(ShapeByType())
83 elif shapes is None:
84 pass
85 else:
86 msg = f"{shapes} is not a valid loman shapes parameter for visualization"
87 raise ValueError(msg)
89 colors = colors.lower()
90 if colors == "state":
91 state_cmap = cmap if isinstance(cmap, dict) else None
92 node_formatters.append(ColorByState(state_cmap)) # type: ignore[arg-type]
93 elif colors == "timing":
94 timing_cmap = cmap if isinstance(cmap, Colormap) else None
95 node_formatters.append(ColorByTiming(timing_cmap))
96 else:
97 msg = f"{colors} is not a valid loman colors parameter for visualization"
98 raise ValueError(msg)
100 node_formatters.append(StandardStylingOverrides())
101 node_formatters.append(RectBlocks())
103 return CompositeNodeFormatter(node_formatters)
106class ColorByState(NodeFormatter):
107 """Node formatter that colors nodes based on their computation state."""
109 DEFAULT_STATE_COLORS: ClassVar[dict[States | None, str]] = {
110 None: "#ffffff", # xkcd white
111 States.PLACEHOLDER: "#f97306", # xkcd orange
112 States.UNINITIALIZED: "#0343df", # xkcd blue
113 States.STALE: "#ffff14", # xkcd yellow
114 States.COMPUTABLE: "#9dff00", # xkcd bright yellow green
115 States.UPTODATE: "#15b01a", # xkcd green
116 States.ERROR: "#e50000", # xkcd red
117 States.PINNED: "#bf77f6", # xkcd light purple
118 }
120 def __init__(self, state_colors: dict[States | None, str] | None = None) -> None:
121 """Initialize with custom state color mapping."""
122 if state_colors is None:
123 state_colors = self.DEFAULT_STATE_COLORS.copy()
124 self.state_colors = state_colors
126 def calibrate(self, nodes: list[Node]) -> None:
127 """Calibrate formatter based on all nodes in the graph."""
128 pass
130 def format(self, name: NodeKey, nodes: list[Node], is_composite: bool) -> dict[str, Any] | None:
131 """Format node color based on computation state."""
132 states = [node.data.get(NodeAttributes.STATE, None) for node in nodes]
133 state: States | None
134 if len(nodes) == 1:
135 state = states[0]
136 else:
137 if any(s == States.ERROR for s in states):
138 state = States.ERROR
139 elif any(s == States.STALE for s in states):
140 state = States.STALE
141 else:
142 state0 = states[0]
143 state = state0 if all(s == state0 for s in states) else None
144 return {"style": "filled", "fillcolor": self.state_colors[state]}
147class ColorByTiming(NodeFormatter):
148 """Node formatter that colors nodes based on their execution timing."""
150 def __init__(self, cmap: Colormap | None = None) -> None:
151 """Initialize with an optional colormap for timing visualization."""
152 if cmap is None:
153 cmap = mpl.colors.LinearSegmentedColormap.from_list("blend", ["#15b01a", "#ffff14", "#e50000"])
154 self.cmap = cmap
155 self.min_duration: float = float("nan")
156 self.max_duration: float = float("nan")
158 def calibrate(self, nodes: list[Node]) -> None:
159 """Calibrate the color mapping based on node timing data."""
160 durations: list[float] = []
161 for node in nodes:
162 timing = node.data.get(NodeAttributes.TIMING)
163 if timing is not None:
164 durations.append(timing.duration)
165 if durations:
166 self.max_duration = max(durations)
167 self.min_duration = min(durations)
169 def format(self, name: NodeKey, nodes: list[Node], is_composite: bool) -> dict[str, Any] | None:
170 """Format a node with timing-based coloring."""
171 if len(nodes) == 1:
172 data = nodes[0].data
173 timing_data = data.get(NodeAttributes.TIMING)
174 if timing_data is None:
175 col = "#FFFFFF"
176 else:
177 duration = timing_data.duration
178 norm_duration: float = (duration - self.min_duration) / max(1e-8, self.max_duration - self.min_duration)
179 col = mpl.colors.rgb2hex(self.cmap(norm_duration))
180 return {"style": "filled", "fillcolor": col}
181 return None
184class ShapeByType(NodeFormatter):
185 """Node formatter that sets node shapes based on their type."""
187 def calibrate(self, nodes: list[Node]) -> None:
188 """Calibrate formatter based on all nodes in the graph."""
189 pass
191 def format(self, name: NodeKey, nodes: list[Node], is_composite: bool) -> dict[str, Any] | None:
192 """Format a node with type-based shape styling."""
193 if len(nodes) == 1:
194 data = nodes[0].data
195 value = data.get(NodeAttributes.VALUE)
196 if value is None:
197 return None
198 if isinstance(value, np.ndarray):
199 return {"shape": "rect"}
200 elif isinstance(value, pd.DataFrame):
201 return {"shape": "box3d"}
202 elif np.isscalar(value):
203 return {"shape": "ellipse"}
204 elif isinstance(value, (list, tuple)):
205 return {"shape": "ellipse", "peripheries": 2}
206 elif isinstance(value, dict):
207 return {"shape": "house", "peripheries": 2}
208 elif isinstance(value, _ComputationLike):
209 return {"shape": "hexagon"}
210 else:
211 return {"shape": "diamond"}
212 return None
215class RectBlocks(NodeFormatter):
216 """Node formatter that shapes composite nodes as rectangles."""
218 def calibrate(self, nodes: list[Node]) -> None:
219 """Calibrate formatter based on all nodes in the graph."""
220 pass
222 def format(self, name: NodeKey, nodes: list[Node], is_composite: bool) -> dict[str, Any] | None:
223 """Return rectangle shape for composite nodes."""
224 if is_composite:
225 return {"shape": "rect", "peripheries": 2}
226 return None
229class StandardLabel(NodeFormatter):
230 """Node formatter that sets node labels."""
232 def calibrate(self, nodes: list[Node]) -> None:
233 """Calibrate formatter based on all nodes in the graph."""
234 pass
236 def format(self, name: NodeKey, nodes: list[Node], is_composite: bool) -> dict[str, Any] | None:
237 """Return standard label for node."""
238 return {"label": name.label}
241def get_group_path(name: NodeKey, data: dict[str, Any]) -> NodeKey:
242 """Determine the group path for a node based on name hierarchy and group attribute."""
243 name_group_path = name.parent
244 attribute_group = data.get(NodeAttributes.GROUP)
245 attribute_group_path = None if attribute_group is None else NodeKey((attribute_group,))
247 group_path = name_group_path.join(attribute_group_path)
248 return group_path
251class StandardGroup(NodeFormatter):
252 """Node formatter that applies standard grouping styles."""
254 def calibrate(self, nodes: list[Node]) -> None:
255 """Calibrate formatter based on all nodes in the graph."""
256 pass
258 def format(self, name: NodeKey, nodes: list[Node], is_composite: bool) -> dict[str, Any] | None:
259 """Format a node with standard group styling."""
260 if len(nodes) == 1:
261 data = nodes[0].data
262 group_path = get_group_path(name, data)
263 else:
264 group_path = name.parent
265 if group_path.is_root:
266 return None
267 return {"_group": group_path}
270class StandardStylingOverrides(NodeFormatter):
271 """Node formatter that applies standard styling overrides."""
273 def calibrate(self, nodes: list[Node]) -> None:
274 """Calibrate formatter based on all nodes in the graph."""
275 pass
277 def format(self, name: NodeKey, nodes: list[Node], is_composite: bool) -> dict[str, Any] | None:
278 """Format a node with standard styling overrides."""
279 if len(nodes) == 1:
280 data = nodes[0].data
281 style = data.get(NodeAttributes.STYLE)
282 if style is None:
283 return None
284 if style == "small":
285 return {"width": 0.3, "height": 0.2, "fontsize": 8}
286 elif style == "dot":
287 return {"shape": "point", "width": 0.1, "peripheries": 1}
288 return None
291@dataclass
292class CompositeNodeFormatter(NodeFormatter):
293 """A node formatter that combines multiple formatters together."""
295 formatters: list[NodeFormatter] = field(default_factory=list)
297 def calibrate(self, nodes: list[Node]) -> None:
298 """Calibrate all the contained formatters with the given nodes."""
299 for formatter in self.formatters:
300 formatter.calibrate(nodes)
302 def format(self, name: NodeKey, nodes: list[Node], is_composite: bool) -> dict[str, Any] | None:
303 """Format a node by combining output from all contained formatters."""
304 d: dict[str, Any] = {}
305 for formatter in self.formatters:
306 format_attrs = formatter.format(name, nodes, is_composite)
307 if format_attrs is not None:
308 d.update(format_attrs)
309 return d
312@dataclass
313class GraphView:
314 """A view for visualizing computation graphs as graphical diagrams."""
316 computation: "Computation"
317 root: Name | None = None
318 node_formatter: NodeFormatter | None = None
319 node_transformations: dict[Name, str] | None = None
320 collapse_all: bool = True
322 graph_attr: dict[str, Any] | None = None
323 node_attr: dict[str, Any] | None = None
324 edge_attr: dict[str, Any] | None = None
326 struct_dag: nx.DiGraph | None = None
327 viz_dag: nx.DiGraph | None = None
328 viz_dot: pydotplus.Dot | None = None
330 def __post_init__(self) -> None:
331 """Initialize the graph view after dataclass construction."""
332 self.refresh()
334 @staticmethod
335 def get_sub_block(
336 dag: nx.DiGraph, root: Name | None, node_transformations: dict[NodeKey, str]
337 ) -> tuple[nx.DiGraph, defaultdict[NodeKey, list[NodeKey]], set[NodeKey]]:
338 """Extract a subgraph with node transformations for visualization."""
339 d_transform_to_nodes: defaultdict[str, list[NodeKey]] = defaultdict(list)
340 for nk, transform in node_transformations.items():
341 d_transform_to_nodes[transform].append(nk)
343 dag_out: nx.DiGraph = nx.DiGraph()
345 d_original_to_mapped: dict[NodeKey, NodeKey] = {}
346 s_collapsed: set[NodeKey] = set()
348 for nk_original in dag.nodes():
349 nk_mapped = nk_original.drop_root(root)
350 if nk_mapped is None:
351 continue
352 nk_highest_collapse = nk_original
353 is_collapsed = False
354 for nk_collapse in d_transform_to_nodes[NodeTransformations.COLLAPSE]:
355 if nk_highest_collapse.is_descendent_of(nk_collapse):
356 nk_highest_collapse = nk_collapse
357 is_collapsed = True
358 nk_mapped = nk_highest_collapse.drop_root(root)
359 if nk_mapped is None: # pragma: no cover
360 continue
361 d_original_to_mapped[nk_original] = nk_mapped
362 if is_collapsed:
363 s_collapsed.add(nk_mapped)
365 for nk_mapped in d_original_to_mapped.values():
366 dag_out.add_node(nk_mapped)
368 for nk_u, nk_v in dag.edges():
369 nk_mapped_u = d_original_to_mapped.get(nk_u)
370 nk_mapped_v = d_original_to_mapped.get(nk_v)
371 if nk_mapped_u is None or nk_mapped_v is None or nk_mapped_u == nk_mapped_v:
372 continue
373 dag_out.add_edge(nk_mapped_u, nk_mapped_v)
375 for nk in d_transform_to_nodes[NodeTransformations.CONTRACT]:
376 contract_node(dag_out, d_original_to_mapped[nk])
377 del d_original_to_mapped[nk]
379 d_mapped_to_original: defaultdict[NodeKey, list[NodeKey]] = defaultdict(list)
380 for nk_original, nk_mapped in d_original_to_mapped.items():
381 if nk_mapped in dag_out.nodes:
382 d_mapped_to_original[nk_mapped].append(nk_original)
384 s_collapsed.intersection_update(dag_out.nodes)
386 return dag_out, d_mapped_to_original, s_collapsed
388 def _initialize_transforms(self) -> dict[NodeKey, str]:
389 """Initialize node transformations for visualization."""
390 node_transformations: dict[NodeKey, str] = {}
391 if self.collapse_all:
392 self._apply_default_collapse_transforms(node_transformations)
393 self._apply_custom_transforms(node_transformations)
394 return node_transformations
396 def _apply_default_collapse_transforms(self, node_transformations: dict[NodeKey, str]) -> None:
397 """Apply default collapse transformations to tree nodes."""
398 for n in self.computation.get_tree_descendents(self.root):
399 nk = to_nodekey(n)
400 if not self.computation.has_node(nk):
401 node_transformations[nk] = NodeTransformations.COLLAPSE
403 def _apply_custom_transforms(self, node_transformations: dict[NodeKey, str]) -> dict[NodeKey, str]:
404 """Apply user-specified custom transformations to nodes."""
405 if self.node_transformations is not None:
406 for rule_name, transform in self.node_transformations.items():
407 include_ancestors = transform == NodeTransformations.EXPAND
408 rule_nk = to_nodekey(rule_name)
409 if is_pattern(rule_nk):
410 apply_nodes: set[NodeKey] = set()
411 for n in self.computation.get_tree_descendents(self.root):
412 nk = to_nodekey(n)
413 if match_pattern(rule_nk, nk):
414 apply_nodes.add(nk)
415 else:
416 apply_nodes = {rule_nk}
417 node_transformations[rule_nk] = transform
418 if include_ancestors:
419 for nk in apply_nodes:
420 for nk1 in nk.ancestors():
421 if nk1.is_root or nk1 == self.root:
422 break
423 node_transformations[nk1] = NodeTransformations.EXPAND
424 for r_nk in apply_nodes:
425 node_transformations[r_nk] = transform
426 return node_transformations
428 def _create_visualization_dag(
429 self, original_nodes: defaultdict[NodeKey, list[NodeKey]], composite_nodes: set[NodeKey]
430 ) -> nx.DiGraph:
431 """Create the visualization DAG from structure and node data."""
432 node_formatter = self.node_formatter
433 if node_formatter is None:
434 node_formatter = NodeFormatter.create()
435 assert self.struct_dag is not None # noqa: S101
436 return create_viz_dag(self.struct_dag, self.computation.dag, node_formatter, original_nodes, composite_nodes)
438 def _create_dot_graph(self) -> pydotplus.Dot:
439 """Create a PyDot graph from the visualization DAG."""
440 return to_pydot(self.viz_dag, self.graph_attr, self.node_attr, self.edge_attr)
442 def refresh(self) -> None:
443 """Refresh the visualization by rebuilding the graph structure."""
444 node_transformations = self._initialize_transforms()
445 self.struct_dag, original_nodes, composite_nodes = self.get_sub_block(
446 self.computation.dag, self.root, node_transformations
447 )
448 self.viz_dag = self._create_visualization_dag(original_nodes, composite_nodes)
449 self.viz_dot = self._create_dot_graph()
451 def svg(self) -> str | None:
452 """Generate SVG representation of the visualization."""
453 if self.viz_dot is None:
454 return None
455 svg_bytes: bytes = self.viz_dot.create_svg() # type: ignore[attr-defined]
456 return svg_bytes.decode("utf-8")
458 def view(self) -> None: # pragma: no cover
459 """Open the visualization in a PDF viewer."""
460 assert self.viz_dot is not None # noqa: S101
461 with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as f:
462 f.write(self.viz_dot.create_pdf()) # type: ignore[attr-defined]
463 if sys.platform == "win32":
464 os.startfile(f.name) # pragma: no cover # nosec B606 # noqa: S606
465 else:
466 subprocess.run(["open", f.name], check=False) # pragma: no cover # nosec B603 B607 # noqa: S603, S607
468 def _repr_svg_(self) -> str | None:
469 """Return SVG representation for Jupyter notebook display."""
470 return self.svg()
473def create_viz_dag(
474 struct_dag: nx.DiGraph,
475 comp_dag: nx.DiGraph,
476 node_formatter: NodeFormatter,
477 original_nodes: defaultdict[NodeKey, list[NodeKey]],
478 composite_nodes: set[NodeKey],
479) -> nx.DiGraph:
480 """Create a visualization DAG from the computation structure."""
481 if node_formatter is not None:
482 nodes: list[Node] = []
483 for nodekey in struct_dag.nodes:
484 for original_nodekey in original_nodes[nodekey]:
485 data = comp_dag.nodes[original_nodekey]
486 n = Node(nodekey, original_nodekey, data)
487 nodes.append(n)
488 node_formatter.calibrate(nodes)
490 viz_dag: nx.DiGraph = nx.DiGraph()
491 node_index_map: dict[NodeKey, str] = {}
492 for i, nodekey in enumerate(struct_dag.nodes):
493 short_name = f"n{i}"
494 attr_dict: dict[str, Any] | None = None
496 if node_formatter is not None:
497 nodes = []
498 for original_nodekey in original_nodes[nodekey]:
499 data = comp_dag.nodes[original_nodekey]
500 n = Node(nodekey, original_nodekey, data)
501 nodes.append(n)
502 is_composite = nodekey in composite_nodes
503 attr_dict = node_formatter.format(nodekey, nodes, is_composite)
504 if attr_dict is None: # pragma: no cover
505 attr_dict = {}
507 attr_dict = {k: v for k, v in attr_dict.items() if v is not None}
509 viz_dag.add_node(short_name, **attr_dict)
510 node_index_map[nodekey] = short_name
512 for name1, name2 in struct_dag.edges():
513 short_name_1 = node_index_map[name1]
514 short_name_2 = node_index_map[name2]
516 group_path1 = get_group_path(name1, struct_dag.nodes[name1])
517 group_path2 = get_group_path(name2, struct_dag.nodes[name2])
518 group_path = NodeKey.common_parent(group_path1, group_path2)
520 edge_attr_dict: dict[str, Any] = {}
521 if not group_path.is_root:
522 # group_path = None
523 edge_attr_dict["_group"] = group_path
525 viz_dag.add_edge(short_name_1, short_name_2, **edge_attr_dict)
527 return viz_dag
530def _group_nodes_and_edges(
531 viz_dag: nx.DiGraph,
532) -> tuple[NodeKey, dict[NodeKey, list[str]], dict[NodeKey, list[tuple[str, str]]]]:
533 """Group nodes and edges by their groups."""
534 root = NodeKey.root()
536 node_groups: dict[NodeKey, list[str]] = {}
537 for name, data in viz_dag.nodes(data=True):
538 group = data.get("_group", root)
539 node_groups.setdefault(group, []).append(name)
541 edge_groups: dict[NodeKey, list[tuple[str, str]]] = {}
542 for name1, name2, data in viz_dag.edges(data=True):
543 group = data.get("_group", root)
544 edge_groups.setdefault(group, []).append((name1, name2))
546 return root, node_groups, edge_groups
549def _create_pydot_nodes(
550 viz_dag: nx.DiGraph,
551 node_groups: dict[NodeKey, list[str]],
552 subgraphs: dict[NodeKey, pydotplus.Dot | pydotplus.Subgraph],
553 root: NodeKey,
554) -> None:
555 """Create PyDot nodes for each group."""
556 for group, names in node_groups.items():
557 c = subgraphs[root] if group is root else create_subgraph(group)
559 for name in names:
560 node = pydotplus.Node(name)
561 for k, v in viz_dag.nodes[name].items():
562 if not k.startswith("_"):
563 node.set(k, v)
564 c.add_node(node)
566 subgraphs[group] = c
569def _ensure_parent_subgraphs(subgraphs: dict[NodeKey, pydotplus.Dot | pydotplus.Subgraph]) -> None:
570 """Ensure all parent subgraphs exist in the hierarchy."""
571 groups = list(subgraphs.keys())
572 for group in groups:
573 group1 = group
574 while True:
575 if group1.is_root:
576 break
577 group1 = group1.parent
578 if group1 in subgraphs:
579 break
580 subgraphs[group1] = create_subgraph(group1)
583def _link_subgraphs(subgraphs: dict[NodeKey, pydotplus.Dot | pydotplus.Subgraph]) -> None:
584 """Link subgraphs to their parents."""
585 for group, subgraph in subgraphs.items():
586 if group.is_root:
587 continue
588 parent = group
589 while True:
590 parent = parent.parent
591 if parent in subgraphs or parent.is_root:
592 break
593 subgraphs[parent].add_subgraph(subgraph)
596def _add_edges_to_subgraphs(
597 edge_groups: dict[NodeKey, list[tuple[str, str]]], subgraphs: dict[NodeKey, pydotplus.Dot | pydotplus.Subgraph]
598) -> None:
599 """Add edges to their respective subgraphs."""
600 for group, edges in edge_groups.items():
601 c = subgraphs[group]
602 for name1, name2 in edges:
603 edge = pydotplus.Edge(name1, name2)
604 c.add_edge(edge)
607def to_pydot(
608 viz_dag: nx.DiGraph | None,
609 graph_attr: dict[str, Any] | None = None,
610 node_attr: dict[str, Any] | None = None,
611 edge_attr: dict[str, Any] | None = None,
612) -> pydotplus.Dot:
613 """Convert a visualization DAG to a PyDot graph for rendering."""
614 assert viz_dag is not None # noqa: S101
615 root, node_groups, edge_groups = _group_nodes_and_edges(viz_dag)
617 subgraphs: dict[NodeKey, pydotplus.Dot | pydotplus.Subgraph] = {
618 root: create_root_graph(graph_attr, node_attr, edge_attr)
619 }
621 _create_pydot_nodes(viz_dag, node_groups, subgraphs, root)
622 _ensure_parent_subgraphs(subgraphs)
623 _link_subgraphs(subgraphs)
624 _add_edges_to_subgraphs(edge_groups, subgraphs)
626 result = subgraphs[root]
627 assert isinstance(result, pydotplus.Dot) # noqa: S101
628 return result
631def create_root_graph(
632 graph_attr: dict[str, Any] | None, node_attr: dict[str, Any] | None, edge_attr: dict[str, Any] | None
633) -> pydotplus.Dot:
634 """Create root Graphviz graph with specified attributes.
636 Notes:
637 Graphviz attributes like size expect a quoted string when containing
638 commas (e.g. "10,8"). Some pydotplus setters don't auto-quote, which
639 can produce a DOT syntax error near ',' if we pass a raw string.
640 We defensively quote string values that contain commas or whitespace.
641 """
643 def _normalize_attr_value(v: Any) -> Any:
644 """Normalize attribute values for Graphviz, quoting strings as needed."""
645 # Keep numeric values as-is
646 if isinstance(v, (int, float)):
647 return v
648 s = str(v)
649 # If already quoted, keep
650 if len(s) >= 2 and ((s[0] == '"' and s[-1] == '"') or (s[0] == "'" and s[-1] == "'")):
651 return s
652 # Quote if contains comma, whitespace, or special characters
653 if any(c in s for c in [",", " ", "\t", "\n"]) or s == "":
654 return f'"{s}"'
655 return s
657 root_graph = pydotplus.Dot()
658 if graph_attr is not None:
659 for k, v in graph_attr.items():
660 root_graph.set(k, _normalize_attr_value(v))
661 if node_attr is not None:
662 # For node/edge defaults, normalize each value too
663 node_defaults = {k: _normalize_attr_value(v) for k, v in node_attr.items()}
664 root_graph.set_node_defaults(**node_defaults)
665 if edge_attr is not None:
666 edge_defaults = {k: _normalize_attr_value(v) for k, v in edge_attr.items()}
667 root_graph.set_edge_defaults(**edge_defaults)
668 return root_graph
671def create_subgraph(group: NodeKey) -> pydotplus.Subgraph:
672 """Create a Graphviz subgraph for a node group."""
673 c = pydotplus.Subgraph("cluster_" + str(group))
674 c.obj_dict["attributes"]["label"] = str(group)
675 return c