Coverage for src / graphable / views / plantuml.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 PlantUmlStylingConfig: 

17 """ 

18 Configuration for customizing PlantUML diagram generation. 

19 

20 Attributes: 

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

22 node_label_fnc: Function to generate the node label. 

23 node_type: PlantUML node type (e.g., 'node', 'component', 'artifact'). 

24 direction: Diagram direction (e.g., 'top to bottom direction', 'left to right direction'). 

25 """ 

26 

27 node_ref_fnc: Callable[[Graphable[Any]], str] = lambda n: f"node_{id(n)}" 

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

29 node_type: str = "node" 

30 direction: str = "top to bottom direction" 

31 cluster_by_tag: bool = False 

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

33 

34 

35def _check_plantuml_on_path() -> None: 

36 """Check if 'plantuml' executable is available in the system path.""" 

37 if which("plantuml") is None: 

38 logger.error("plantuml not found on PATH.") 

39 raise FileNotFoundError( 

40 "plantuml is required but not available on $PATH. " 

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

42 ) 

43 

44 

45def create_topology_plantuml( 

46 graph: Graph, config: PlantUmlStylingConfig | None = None 

47) -> str: 

48 """ 

49 Generate PlantUML definition from a Graph. 

50 

51 Args: 

52 graph (Graph): The graph to convert. 

53 config (PlantUmlStylingConfig | None): Styling configuration. 

54 

55 Returns: 

56 str: The PlantUML definition string. 

57 """ 

58 config = config or PlantUmlStylingConfig() 

59 puml: list[str] = ["@startuml", f"{config.direction}"] 

60 

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

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

63 return None 

64 sorted_tags = config.tag_sort_fnc(node.tags) 

65 return sorted_tags[0] if sorted_tags else None 

66 

67 # Group nodes by cluster 

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

69 for node in graph.topological_order(): 

70 cluster = get_cluster(node) 

71 if cluster not in clusters: 

72 clusters[cluster] = [] 

73 clusters[cluster].append(node) 

74 

75 # Nodes (potentially within clusters) 

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

77 indent = "" 

78 if cluster_name: 

79 puml.append(f'package "{cluster_name}" {{') 

80 indent = " " 

81 

82 for node in nodes: 

83 node_ref = config.node_ref_fnc(node) 

84 node_label = config.node_label_fnc(node) 

85 puml.append(f'{indent}{config.node_type} "{node_label}" as {node_ref}') 

86 

87 if cluster_name: 

88 puml.append("}") 

89 

90 # Edges 

91 for node in graph.topological_order(): 

92 node_ref = config.node_ref_fnc(node) 

93 for dependent, _ in graph.internal_dependents(node): 

94 dep_ref = config.node_ref_fnc(dependent) 

95 puml.append(f"{node_ref} --> {dep_ref}") 

96 

97 puml.append("@enduml") 

98 return "\n".join(puml) 

99 

100 

101@register_view(".puml", creator_fnc=create_topology_plantuml) 

102def export_topology_plantuml( 

103 graph: Graph, output: Path, config: PlantUmlStylingConfig | None = None 

104) -> None: 

105 """ 

106 Export the graph to a PlantUML (.puml) file. 

107 

108 Args: 

109 graph (Graph): The graph to export. 

110 output (Path): The output file path. 

111 config (PlantUMLStylingConfig | None): Styling configuration. 

112 """ 

113 logger.info(f"Exporting PlantUML definition to: {output}") 

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

115 f.write(create_topology_plantuml(graph, config)) 

116 

117 

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

119def export_topology_plantuml_image( 

120 graph: Graph, output: Path, config: PlantUmlStylingConfig | None = None 

121) -> None: 

122 """ 

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

124 

125 Args: 

126 graph (Graph): The graph to export. 

127 output (Path): The output file path. 

128 config (PlantUmlStylingConfig | None): Styling configuration. 

129 """ 

130 logger.info(f"Exporting PlantUML image to: {output}") 

131 _check_plantuml_on_path() 

132 

133 p = Path(output) 

134 puml_content: str = create_topology_plantuml(graph, config) 

135 

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

137 

138 try: 

139 run( 

140 ["plantuml", f"-t{fmt}", "-p"], 

141 input=puml_content, 

142 check=True, 

143 stderr=PIPE, 

144 stdout=open(p, "wb"), 

145 text=True, 

146 ) 

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

148 except CalledProcessError as e: 

149 logger.error(f"Error executing plantuml: {e.stderr}") 

150 raise 

151 except Exception as e: 

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

153 raise