diff --git a/codemap/core/mermaid_generator.py b/codemap/core/mermaid_generator.py new file mode 100644 index 0000000..d8d416a --- /dev/null +++ b/codemap/core/mermaid_generator.py @@ -0,0 +1,98 @@ +from typing import Dict, List, Optional +from codemap.core.graph_builder import GraphData, Node, Edge + + +STYLE_COLORS = { + "python": ("#3572A5", "#FFE873"), + "javascript": ("#F1E05A", "#F7DF9E"), + "go": ("#00ADD8", "#A8D7EC"), + "default": ("#6e7681", "#e1e4e8") +} + + +class MermaidGenerator: + def __init__(self, graph_data: GraphData): + self.graph_data = graph_data + + def generate_flowchart(self, include_packages: bool = True) -> str: + lines = ["graph TD"] + lines.append(" classDef python fill:#3572A5,stroke:#333,stroke-width:2px,color:white") + lines.append(" classDef javascript fill:#F1E05A,stroke:#333,stroke-width:2px,color:black") + lines.append(" classDef go fill:#00ADD8,stroke:#333,stroke-width:2px,color:white") + lines.append(" classDef default fill:#6e7681,stroke:#333,stroke-width:2px,color:white") + + if include_packages and self.graph_data.packages: + for package, nodes in self.graph_data.packages.items(): + lines.append(f" subgraph {package}") + for node_id in nodes: + node = self._find_node(node_id) + if node: + safe_id = self._safe_id(node_id) + lines.append(f" {safe_id}[{node.name}]") + lines.append(" end") + + node_ids: List[str] = [] + for node in self.graph_data.nodes: + safe_id = self._safe_id(node.id) + node_ids.append(safe_id) + fill_color, _ = STYLE_COLORS.get(node.file_type, STYLE_COLORS["default"]) + style = f" {safe_id}[{node.name}]:::python" if node.file_type == "python" else \ + f" {safe_id}[{node.name}]:::javascript" if node.file_type == "javascript" else \ + f" {safe_id}[{node.name}]:::go" if node.file_type == "go" else \ + f" {safe_id}[{node.name}]:::default" + lines.append(style) + + for edge in self.graph_data.edges: + if self._find_node(edge.source) and self._find_node(edge.target): + source_id = self._safe_id(edge.source) + target_id = self._safe_id(edge.target) + if edge.label: + lines.append(f" {source_id} -->|{edge.label}| {target_id}") + else: + lines.append(f" {source_id} --> {target_id}") + + return "\n".join(lines) + + def _find_node(self, node_id: str) -> Optional[Node]: + for node in self.graph_data.nodes: + if node.id == node_id: + return node + return None + + def _safe_id(self, node_id: str) -> str: + safe = node_id.replace("/", "_").replace(".", "_").replace("-", "_") + if safe[0].isdigit(): + safe = "n" + safe + return safe[:50] + + def generate_class_diagram(self) -> str: + lines = ["classDiagram"] + lines.append(" class Node {") + lines.append(" +str id") + lines.append(" +str name") + lines.append(" +str file_type") + lines.append(" }") + lines.append(" class Edge {") + lines.append(" +str source") + lines.append(" +str target") + lines.append(" +str label") + lines.append(" }") + + for node in self.graph_data.nodes: + lines.append(f" Node <|-- {node.name}") + + for edge in self.graph_data.edges: + lines.append(f" {edge.source} --> {edge.target} : {edge.label}") + + return "\n".join(lines) + + def generate_pie_chart(self) -> str: + type_counts: Dict[str, int] = {} + for node in self.graph_data.nodes: + type_counts[node.file_type] = type_counts.get(node.file_type, 0) + 1 + + lines = ["pie title File Types"] + for file_type, count in type_counts.items(): + lines.append(f' "{file_type}" : {count}') + + return "\n".join(lines)