diff --git a/src/exporters/dot.py b/src/exporters/dot.py new file mode 100644 index 0000000..b6824fc --- /dev/null +++ b/src/exporters/dot.py @@ -0,0 +1,71 @@ +from pathlib import Path + +from src.graph.builder import GraphBuilder + + +class DOTExporter: + def __init__(self, graph_builder: GraphBuilder): + self.graph_builder = graph_builder + + def export(self, output_path: Path, layout: str = "dot", rankdir: str = "TB") -> None: + lines = [] + lines.append("digraph codegraph {") + lines.append(f' rankdir={rankdir};') + lines.append(' node [fontname="Helvetica"];') + lines.append(' edge [fontname="Helvetica"];') + lines.append("") + + for node in self.graph_builder.get_nodes(): + lines.append(f' "{node.node_id}" [') + lines.append(f' label="{node.label}"') + lines.append(f' shape={node.shape}') + lines.append(f' style={node.style}') + lines.append(f' fillcolor="{node.color}"') + lines.append(' color="#333333"') + lines.append(' ];') + + lines.append("") + + for edge in self.graph_builder.edges: + edge_style = "dashed" if edge.edge_type == "calls" else "solid" + lines.append(f' "{edge.source}" -> "{edge.target}" [') + lines.append(f' label="{edge.label}"') + lines.append(f' style={edge_style}') + lines.append(' arrowhead=normal') + lines.append(' ];') + + lines.append("}") + + content = "\n".join(lines) + + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_text(content) + + def get_string(self, layout: str = "dot", rankdir: str = "TB") -> str: + lines = [] + lines.append("digraph codegraph {") + lines.append(f' rankdir={rankdir};') + lines.append(' node [fontname="Helvetica"];') + lines.append(' edge [fontname="Helvetica"];') + lines.append("") + + for node in self.graph_builder.get_nodes(): + lines.append(f' "{node.node_id}" [') + lines.append(f' label="{node.label}"') + lines.append(f' shape={node.shape}') + lines.append(f' style={node.style}') + lines.append(f' fillcolor="{node.color}"') + lines.append(' ];') + + lines.append("") + + for edge in self.graph_builder.edges: + edge_style = "dashed" if edge.edge_type == "calls" else "solid" + lines.append(f' "{edge.source}" -> "{edge.target}" [') + lines.append(f' label="{edge.label}"') + lines.append(f' style={edge_style}') + lines.append(' ];') + + lines.append("}") + + return "\n".join(lines)