Initial commit: CodeMap v0.1.0 - CLI tool for code analysis and diagram generation

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)