ai-coding-assistant 0.5.0__py3-none-any.whl

coding_assistant/documentation/generators/diagram_generator.py

@@ -0,0 +1,511 @@
+"""Generate Mermaid diagrams from code structure.
+
+This module provides generation of various diagram types in Mermaid format:
+- Architecture diagrams (module/partition dependencies)
+- Class diagrams (classes, inheritance, relationships)
+- Sequence diagrams (function call flows)
+- Dependency graphs (file/module dependencies)
+"""
+
+from typing import List, Dict, Set, Optional, Tuple
+from pathlib import Path
+import networkx as nx
+import re
+
+from coding_assistant.documentation.decomposition.partitioner import Partition
+from coding_assistant.utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+class MermaidDiagramGenerator:
+    """
+    Generate various Mermaid diagrams from code structure.
+
+    Supports:
+    - Architecture diagrams: High-level module relationships
+    - Class diagrams: OOP structure and inheritance
+    - Sequence diagrams: Function call sequences
+    - Dependency graphs: File/module dependencies
+    """
+
+    def __init__(self):
+        """Initialize the diagram generator."""
+        pass
+
+    def generate_architecture_diagram(self,
+                                        partitions: List[Partition],
+                                        dependency_graph: Optional[nx.DiGraph] = None,
+                                        direction: str = "TB") -> str:
+        """
+        Generate high-level architecture diagram showing modules and their relationships.
+
+        Args:
+            partitions: List of repository partitions
+            dependency_graph: Optional dependency graph for additional context
+            direction: Graph direction (TB=top-bottom, LR=left-right, RL, BT)
+
+        Returns:
+            Mermaid diagram as string
+        """
+        logger.info(f"Generating architecture diagram for {len(partitions)} partitions")
+
+        mermaid = [f"graph {direction}"]
+
+        # Add nodes for each partition
+        for partition in partitions:
+            node_id = self._sanitize_id(partition.name)
+            label = self._format_partition_label(partition.name, partition.size_loc)
+
+            # Choose node shape based on architectural role (if available)
+            node_shape = self._get_node_shape_for_partition(partition)
+            mermaid.append(f"    {node_id}{node_shape[0]}{label}{node_shape[1]}")
+
+        # Add edges for dependencies
+        added_edges = set()
+        for partition in partitions:
+            node_id = self._sanitize_id(partition.name)
+
+            for dep in partition.dependencies:
+                dep_id = self._sanitize_id(dep)
+                edge_key = (node_id, dep_id)
+
+                # Avoid duplicate edges
+                if edge_key not in added_edges:
+                    mermaid.append(f"    {node_id} --> {dep_id}")
+                    added_edges.add(edge_key)
+
+        # Add styling
+        mermaid.extend(self._add_architecture_styling(partitions))
+
+        result = "\n".join(mermaid)
+        logger.debug(f"Architecture diagram generated ({len(mermaid)} lines)")
+
+        return result
+
+    def generate_class_diagram(self,
+                                 parsed_files: Dict[str, Dict],
+                                 file_filter: Optional[List[str]] = None,
+                                 max_classes: int = 20) -> str:
+        """
+        Generate class diagram showing classes, methods, and inheritance.
+
+        Args:
+            parsed_files: Dictionary of file_path -> parsed code data
+            file_filter: Optional list of files to include (None = all)
+            max_classes: Maximum number of classes to include
+
+        Returns:
+            Mermaid class diagram as string
+        """
+        logger.info(f"Generating class diagram from {len(parsed_files)} files")
+
+        mermaid = ["classDiagram"]
+
+        # Extract all classes from parsed files
+        all_classes = []
+        class_relationships = []
+
+        for file_path, parsed_data in parsed_files.items():
+            # Apply file filter
+            if file_filter and file_path not in file_filter:
+                continue
+
+            classes = parsed_data.get('classes', [])
+
+            for cls in classes:
+                all_classes.append({
+                    'file': file_path,
+                    'data': cls
+                })
+
+                # Track inheritance relationships
+                if cls.get('bases'):
+                    for base in cls['bases']:
+                        class_relationships.append((base, cls['name'], 'inherits'))
+
+        # Limit to max_classes (prioritize classes with relationships)
+        if len(all_classes) > max_classes:
+            # Sort by number of relationships (methods count as a heuristic)
+            all_classes.sort(
+                key=lambda c: len(c['data'].get('methods', [])),
+                reverse=True
+            )
+            all_classes = all_classes[:max_classes]
+
+        # Generate class definitions
+        included_classes = set()
+        for cls_info in all_classes:
+            cls = cls_info['data']
+            class_name = cls['name']
+            included_classes.add(class_name)
+
+            mermaid.append(f"    class {class_name} {{")
+
+            # Add attributes (if available)
+            attributes = cls.get('attributes', [])
+            for attr in attributes[:10]:  # Limit attributes
+                attr_name = attr.get('name', attr) if isinstance(attr, dict) else attr
+                attr_type = attr.get('type', '') if isinstance(attr, dict) else ''
+                if attr_type:
+                    mermaid.append(f"        +{attr_type} {attr_name}")
+                else:
+                    mermaid.append(f"        +{attr_name}")
+
+            # Add methods
+            methods = cls.get('methods', [])
+            for method in methods[:15]:  # Limit methods
+                method_name = method.get('name', method) if isinstance(method, dict) else method
+
+                # Skip private methods starting with _
+                if method_name.startswith('_') and not method_name.startswith('__'):
+                    continue
+
+                # Get method signature if available
+                params = method.get('parameters', []) if isinstance(method, dict) else []
+                return_type = method.get('return_type', '') if isinstance(method, dict) else ''
+
+                if params:
+                    param_str = ', '.join(params[:3])  # Limit params shown
+                    if len(params) > 3:
+                        param_str += ', ...'
+                else:
+                    param_str = ''
+
+                if return_type:
+                    mermaid.append(f"        +{method_name}({param_str}) {return_type}")
+                else:
+                    mermaid.append(f"        +{method_name}({param_str})")
+
+            mermaid.append("    }")
+
+        # Add inheritance relationships
+        for base, derived, rel_type in class_relationships:
+            # Only include if both classes are in the diagram
+            if base in included_classes and derived in included_classes:
+                if rel_type == 'inherits':
+                    mermaid.append(f"    {base} <|-- {derived}")
+
+        result = "\n".join(mermaid)
+        logger.debug(f"Class diagram generated with {len(included_classes)} classes")
+
+        return result
+
+    def generate_sequence_diagram(self,
+                                    function_name: str,
+                                    call_chain: List[Dict],
+                                    max_depth: int = 10) -> str:
+        """
+        Generate sequence diagram for a function's execution flow.
+
+        Args:
+            function_name: Name of the main function
+            call_chain: List of function calls with structure:
+                        [{'caller': 'A', 'callee': 'B', 'message': 'method()'}]
+            max_depth: Maximum call depth to visualize
+
+        Returns:
+            Mermaid sequence diagram as string
+        """
+        logger.info(f"Generating sequence diagram for {function_name}")
+
+        mermaid = ["sequenceDiagram"]
+
+        # Add title
+        mermaid.append(f"    title {function_name} Execution Flow")
+
+        # Extract unique participants
+        participants = self._extract_participants(call_chain)
+
+        # Add participants
+        for participant in participants[:max_depth]:
+            mermaid.append(f"    participant {participant}")
+
+        # Add interactions
+        call_depth = 0
+        for call in call_chain[:max_depth * 2]:  # Limit total calls
+            caller = call.get('caller', 'Unknown')
+            callee = call.get('callee', 'Unknown')
+            message = call.get('message', callee)
+            call_type = call.get('type', 'sync')  # sync, async, return
+
+            # Skip if participants not included
+            if caller not in participants or callee not in participants:
+                continue
+
+            # Different arrow types based on call type
+            if call_type == 'async':
+                mermaid.append(f"    {caller}-))+{callee}: {message}")
+            elif call_type == 'return':
+                mermaid.append(f"    {callee}-->>-{caller}: return")
+                call_depth -= 1
+            else:
+                mermaid.append(f"    {caller}->>+{callee}: {message}")
+                call_depth += 1
+
+            # Prevent excessive nesting
+            if call_depth > max_depth:
+                break
+
+        result = "\n".join(mermaid)
+        logger.debug(f"Sequence diagram generated with {len(participants)} participants")
+
+        return result
+
+    def generate_dependency_graph(self,
+                                    dependency_graph: nx.DiGraph,
+                                    max_nodes: int = 50,
+                                    show_files: bool = True) -> str:
+        """
+        Generate dependency graph visualization.
+
+        Args:
+            dependency_graph: NetworkX directed graph of dependencies
+            max_nodes: Maximum nodes to include (most important)
+            show_files: If True, show file names; otherwise show modules
+
+        Returns:
+            Mermaid dependency graph as string
+        """
+        logger.info(f"Generating dependency graph with {dependency_graph.number_of_nodes()} nodes")
+
+        # Limit to most important nodes if graph is too large
+        if dependency_graph.number_of_nodes() > max_nodes:
+            important_nodes = self._get_important_nodes(dependency_graph, max_nodes)
+            subgraph = dependency_graph.subgraph(important_nodes).copy()
+        else:
+            subgraph = dependency_graph
+
+        mermaid = ["graph LR"]
+
+        # Add nodes
+        for node in subgraph.nodes():
+            node_id = self._sanitize_id(node)
+
+            # Format label based on show_files
+            if show_files:
+                label = Path(node).name if '/' in node or '\\' in node else node
+            else:
+                label = node
+
+            # Determine node styling based on centrality
+            node_data = subgraph.nodes.get(node, {})
+            is_central = node_data.get('is_central', False)
+
+            if is_central:
+                mermaid.append(f"    {node_id}[[\"{label}\"]]")  # Double border for central nodes
+            else:
+                mermaid.append(f"    {node_id}[\"{label}\"]")
+
+        # Add edges
+        for source, target in subgraph.edges():
+            source_id = self._sanitize_id(source)
+            target_id = self._sanitize_id(target)
+
+            # Check for edge data (weight, type, etc.)
+            edge_data = subgraph.edges.get((source, target), {})
+            weight = edge_data.get('weight', 1)
+
+            # Thicker arrows for stronger dependencies
+            if weight > 5:
+                mermaid.append(f"    {source_id} ==> {target_id}")
+            else:
+                mermaid.append(f"    {source_id} --> {target_id}")
+
+        # Add styling for central nodes
+        mermaid.append("    classDef central fill:#ffd700,stroke:#ff6347,stroke-width:3px")
+        central_nodes = [
+            self._sanitize_id(n) for n in subgraph.nodes()
+            if subgraph.nodes.get(n, {}).get('is_central', False)
+        ]
+        if central_nodes:
+            mermaid.append(f"    class {','.join(central_nodes)} central")
+
+        result = "\n".join(mermaid)
+        logger.debug(f"Dependency graph generated with {subgraph.number_of_nodes()} nodes")
+
+        return result
+
+    # Helper methods
+
+    def _sanitize_id(self, name: str) -> str:
+        """
+        Sanitize name for use as Mermaid ID.
+
+        Creates clean, valid Mermaid node IDs:
+        - Uses basename for paths (shorter, more readable)
+        - Replaces special characters with underscores
+        - Ensures ID starts with a letter
+        - Collapses multiple underscores
+        """
+        # Extract basename if it's a path (use last meaningful component)
+        if '/' in name or '\\' in name:
+            # Get the last non-empty path component
+            parts = re.split(r'[/\\]', name)
+            meaningful_parts = [p for p in parts if p and p not in ('.', '..')]
+            if meaningful_parts:
+                # Use last 2 parts for context (e.g., "module/file" instead of just "file")
+                name = '_'.join(meaningful_parts[-2:]) if len(meaningful_parts) > 1 else meaningful_parts[-1]
+
+        # Replace path separators, dots, dashes, spaces with underscores
+        sanitized = re.sub(r'[/\\.\-\s]', '_', name)
+
+        # Remove any remaining special characters (keep alphanumeric and underscore)
+        sanitized = re.sub(r'[^a-zA-Z0-9_]', '', sanitized)
+
+        # Collapse multiple consecutive underscores into one
+        sanitized = re.sub(r'_+', '_', sanitized)
+
+        # Remove leading/trailing underscores
+        sanitized = sanitized.strip('_')
+
+        # Ensure ID starts with a letter (Mermaid requirement)
+        if sanitized and not sanitized[0].isalpha():
+            sanitized = 'n_' + sanitized
+
+        # Handle empty result
+        if not sanitized:
+            sanitized = 'unknown'
+
+        return sanitized
+
+    def _format_partition_label(self, name: str, size_loc: int) -> str:
+        """Format partition node label with size info."""
+        # Extract meaningful name from path
+        if '/' in name or '\\' in name:
+            parts = re.split(r'[/\\]', name)
+            meaningful_parts = [p for p in parts if p and p not in ('.', '..', '')]
+            if meaningful_parts:
+                # Use last 2 meaningful parts for context
+                name = '/'.join(meaningful_parts[-2:]) if len(meaningful_parts) > 1 else meaningful_parts[-1]
+
+        # Remove leading dots/special chars from name
+        name = name.lstrip('./\\')
+
+        # Shorten long names
+        if len(name) > 30:
+            name = name[:27] + '...'
+
+        # Escape quotes in name
+        name = name.replace('"', "'")
+
+        return f'"{name}\\n({size_loc} LOC)"'
+
+    def _get_node_shape_for_partition(self, partition: Partition) -> Tuple[str, str]:
+        """
+        Determine node shape based on partition characteristics.
+
+        Returns:
+            Tuple of (opening bracket, closing bracket) for node shape
+        """
+        # Check if partition has metadata with architectural role
+        role = partition.metadata.get('architectural_role', '').lower()
+
+        if 'api' in role or 'interface' in role:
+            return ('([', '])')  # Stadium shape for APIs
+        elif 'data' in role or 'model' in role:
+            return ('[(', ')]')  # Cylindrical shape for data
+        elif 'ui' in role or 'view' in role:
+            return ('{{', '}}')  # Hexagon for UI
+        else:
+            return ('[', ']')    # Rectangle for general
+
+    def _add_architecture_styling(self, partitions: List[Partition]) -> List[str]:
+        """Add styling for architecture diagram."""
+        styles = []
+
+        # Define color scheme based on size
+        large_partitions = []
+        medium_partitions = []
+        small_partitions = []
+
+        for partition in partitions:
+            node_id = self._sanitize_id(partition.name)
+
+            if partition.size_loc > 2000:
+                large_partitions.append(node_id)
+            elif partition.size_loc > 500:
+                medium_partitions.append(node_id)
+            else:
+                small_partitions.append(node_id)
+
+        # Define class styles
+        styles.append("    classDef large fill:#ff6b6b,stroke:#c92a2a,stroke-width:2px")
+        styles.append("    classDef medium fill:#4dabf7,stroke:#1971c2,stroke-width:2px")
+        styles.append("    classDef small fill:#51cf66,stroke:#2f9e44,stroke-width:2px")
+
+        # Apply styles
+        if large_partitions:
+            styles.append(f"    class {','.join(large_partitions)} large")
+        if medium_partitions:
+            styles.append(f"    class {','.join(medium_partitions)} medium")
+        if small_partitions:
+            styles.append(f"    class {','.join(small_partitions)} small")
+
+        return styles
+
+    def _extract_participants(self, call_chain: List[Dict]) -> List[str]:
+        """Extract unique participants from call chain."""
+        participants = []
+        seen = set()
+
+        for call in call_chain:
+            caller = call.get('caller', 'Unknown')
+            callee = call.get('callee', 'Unknown')
+
+            if caller not in seen:
+                participants.append(caller)
+                seen.add(caller)
+
+            if callee not in seen:
+                participants.append(callee)
+                seen.add(callee)
+
+        return participants
+
+    def _get_important_nodes(self,
+                              graph: nx.DiGraph,
+                              max_nodes: int) -> List[str]:
+        """
+        Get most important nodes from graph based on centrality metrics.
+
+        Uses PageRank and degree centrality to identify key nodes.
+        """
+        try:
+            # Calculate PageRank
+            pagerank = nx.pagerank(graph)
+
+            # Calculate degree centrality
+            degree_centrality = nx.degree_centrality(graph)
+
+            # Combine scores (weighted average)
+            combined_scores = {}
+            for node in graph.nodes():
+                combined_scores[node] = (
+                    0.7 * pagerank.get(node, 0) +
+                    0.3 * degree_centrality.get(node, 0)
+                )
+
+            # Sort by score and take top max_nodes
+            sorted_nodes = sorted(
+                combined_scores.items(),
+                key=lambda x: x[1],
+                reverse=True
+            )
+
+            important_nodes = [node for node, score in sorted_nodes[:max_nodes]]
+
+            # Mark central nodes in graph
+            for node in important_nodes[:max_nodes // 5]:  # Top 20% are "central"
+                if node in graph.nodes():
+                    graph.nodes[node]['is_central'] = True
+
+            return important_nodes
+
+        except Exception as e:
+            logger.warning(f"Error calculating node importance: {e}, using degree instead")
+
+            # Fallback: just use degree
+            degree_dict = dict(graph.degree())
+            sorted_nodes = sorted(degree_dict.items(), key=lambda x: x[1], reverse=True)
+            return [node for node, degree in sorted_nodes[:max_nodes]]

coding_assistant/documentation/graph/__init__.py

@@ -0,0 +1,13 @@
+"""Graph analysis for code dependencies and module relationships."""
+
+from coding_assistant.documentation.graph.dependency_builder import (
+    DependencyGraphBuilder,
+    CodeEntity,
+)
+from coding_assistant.documentation.graph.module_analyzer import ModuleAnalyzer
+
+__all__ = [
+    'DependencyGraphBuilder',
+    'CodeEntity',
+    'ModuleAnalyzer',
+]