Coverage for src / graphable / views / graphviz.py: 100%

1from dataclasses import dataclass 

2from logging import getLogger 

3from pathlib import Path 

4from shutil import which 

5from subprocess import PIPE, CalledProcessError, run 

6from typing import Any, Callable 

7 

8from ..graph import Graph 

9from ..graphable import Graphable 

10from ..registry import register_view 

11 

12logger = getLogger(__name__) 

13 

14 

15@dataclass 

16class GraphvizStylingConfig: 

17 """ 

18 Configuration for customizing Graphviz DOT generation. 

19 

20 Attributes: 

21 node_ref_fnc: Function to generate the node identifier (reference). 

22 node_label_fnc: Function to generate the node label. 

23 node_attr_fnc: Function to generate a dictionary of attributes for a node. 

24 edge_attr_fnc: Function to generate a dictionary of attributes for an edge. 

25 graph_attr: Dictionary of global graph attributes. 

26 node_attr_default: Dictionary of default node attributes. 

27 edge_attr_default: Dictionary of default edge attributes. 

28 """ 

29 

30 node_ref_fnc: Callable[[Graphable[Any]], str] = lambda n: str(n.reference) 

31 node_label_fnc: Callable[[Graphable[Any]], str] = lambda n: str(n.reference) 

32 node_attr_fnc: Callable[[Graphable[Any]], dict[str, str]] | None = None 

33 edge_attr_fnc: Callable[[Graphable[Any], Graphable[Any]], dict[str, str]] | None = ( 

34 None 

35 ) 

36 graph_attr: dict[str, str] | None = None 

37 node_attr_default: dict[str, str] | None = None 

38 edge_attr_default: dict[str, str] | None = None 

39 cluster_by_tag: bool = False 

40 tag_sort_fnc: Callable[[set[str]], list[str]] = lambda s: sorted(list(s)) 

41 

42 

43def _check_dot_on_path() -> None: 

44 """Check if 'dot' executable is available in the system path.""" 

45 if which("dot") is None: 

46 logger.error("dot not found on PATH.") 

47 raise FileNotFoundError( 

48 "dot (Graphviz) is required but not available on $PATH. " 

49 "Install it via your package manager (e.g., 'sudo apt install graphviz')." 

50 ) 

51 

52 

53def _escape_dot_string(s: str) -> str: 

54 """Escape double quotes and backslashes for DOT strings.""" 

55 return s.replace("\\", "\\\\").replace('"', '\\"') 

56 

57 

58def _format_attrs(attrs: dict[str, str] | None) -> str: 

59 """Format a dictionary of attributes into a DOT attribute string.""" 

60 if not attrs: 

61 return "" 

62 parts = [f'{k}="{_escape_dot_string(v)}"' for k, v in attrs.items()] 

63 return f" [{', '.join(parts)}]" 

64 

65 

66def create_topology_graphviz_dot( 

67 graph: Graph, config: GraphvizStylingConfig | None = None 

68) -> str: 

69 """ 

70 Generate Graphviz DOT definition from a Graph. 

71 

72 Args: 

73 graph (Graph): The graph to convert. 

74 config (GraphvizStylingConfig | None): Styling configuration. 

75 

76 Returns: 

77 str: The Graphviz DOT definition string. 

78 """ 

79 config = config or GraphvizStylingConfig() 

80 dot: list[str] = ["digraph G {"] 

81 

82 # Global attributes 

83 if config.graph_attr: 

84 for k, v in config.graph_attr.items(): 

85 dot.append(f' {k}="{_escape_dot_string(v)}";') 

86 

87 if config.node_attr_default: 

88 dot.append(f" node{_format_attrs(config.node_attr_default)};") 

89 

90 if config.edge_attr_default: 

91 dot.append(f" edge{_format_attrs(config.edge_attr_default)};") 

92 

93 def get_cluster(node: Graphable[Any]) -> str | None: 

94 if not config.cluster_by_tag or not node.tags: 

95 return None 

96 sorted_tags = config.tag_sort_fnc(node.tags) 

97 return sorted_tags[0] if sorted_tags else None 

98 

99 # Group nodes by cluster 

100 clusters: dict[str | None, list[Graphable[Any]]] = {} 

101 for node in graph.topological_order(): 

102 cluster = get_cluster(node) 

103 if cluster not in clusters: 

104 clusters[cluster] = [] 

105 clusters[cluster].append(node) 

106 

107 # Nodes (potentially within clusters) 

108 for cluster_name, nodes in clusters.items(): 

109 indent = " " 

110 if cluster_name: 

111 safe_cluster = _escape_dot_string(cluster_name) 

112 dot.append(f' subgraph "cluster_{safe_cluster}" {{') 

113 dot.append(f' label="{safe_cluster}";') 

114 indent = " " 

115 

116 for node in nodes: 

117 node_ref = _escape_dot_string(config.node_ref_fnc(node)) 

118 node_attrs = {"label": config.node_label_fnc(node)} 

119 if config.node_attr_fnc: 

120 node_attrs.update(config.node_attr_fnc(node)) 

121 

122 dot.append(f'{indent}"{node_ref}"{_format_attrs(node_attrs)};') 

123 

124 if cluster_name: 

125 dot.append(" }") 

126 

127 # Edges 

128 for node in graph.topological_order(): 

129 node_ref = _escape_dot_string(config.node_ref_fnc(node)) 

130 for dependent, attrs in graph.internal_dependents(node): 

131 dep_ref = _escape_dot_string(config.node_ref_fnc(dependent)) 

132 edge_attrs = {} 

133 if config.edge_attr_fnc: 

134 edge_attrs.update(config.edge_attr_fnc(node, dependent)) 

135 

136 dot.append(f' "{node_ref}" -> "{dep_ref}"{_format_attrs(edge_attrs)};') 

137 

138 dot.append("}") 

139 return "\n".join(dot) 

140 

141 

142@register_view([".dot", ".gv"], creator_fnc=create_topology_graphviz_dot) 

143def export_topology_graphviz_dot( 

144 graph: Graph, output: Path, config: GraphvizStylingConfig | None = None 

145) -> None: 

146 """ 

147 Export the graph to a Graphviz .dot file. 

148 

149 Args: 

150 graph (Graph): The graph to export. 

151 output (Path): The output file path. 

152 config (GraphvizStylingConfig | None): Styling configuration. 

153 """ 

154 logger.info(f"Exporting graphviz dot to: {output}") 

155 with open(output, "w+") as f: 

156 f.write(create_topology_graphviz_dot(graph, config)) 

157 

158 

159@register_view([".svg", ".png"]) 

160def export_topology_graphviz_image( 

161 graph: Graph, output: Path, config: GraphvizStylingConfig | None = None 

162) -> None: 

163 """ 

164 Export the graph to an image file (SVG or PNG) using the 'dot' command. 

165 

166 Args: 

167 graph (Graph): The graph to export. 

168 output (Path): The output file path. 

169 config (GraphvizStylingConfig | None): Styling configuration. 

170 """ 

171 logger.info(f"Exporting graphviz image to: {output}") 

172 _check_dot_on_path() 

173 

174 p = Path(output) 

175 dot_content: str = create_topology_graphviz_dot(graph, config) 

176 

177 # Extension without dot 

178 fmt = p.suffix[1:].lower() 

179 

180 try: 

181 run( 

182 ["dot", f"-T{fmt}", "-o", str(p)], 

183 input=dot_content, 

184 check=True, 

185 stderr=PIPE, 

186 stdout=PIPE, 

187 text=True, 

188 ) 

189 logger.info(f"Successfully exported {fmt.upper()} to {output}") 

190 except CalledProcessError as e: 

191 logger.error(f"Error executing dot: {e.stderr}") 

192 raise 

193 except Exception as e: 

194 logger.error(f"Failed to export {fmt.upper()} to {output}: {e}") 

195 raise