code2llm 0.3.7__py3-none-any.whl

code2flow/refactor/prompt_engine.py

@@ -0,0 +1,150 @@
+"""Engine for generating refactoring prompts using Jinja2 templates."""
+import os
+import jinja2
+import tiktoken
+from tree_sitter import Language, Parser
+import tree_sitter_python
+from typing import List, Dict, Any, Optional
+from ..core.models import AnalysisResult, CodeSmell
+
+class PromptEngine:
+    """Generate refactoring prompts from analysis results and detected smells."""
+    
+    def __init__(self, result: AnalysisResult, template_dir: Optional[str] = None):
+        if template_dir is None:
+            # Default to templates directory relative to this file
+            template_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), 'templates')
+            
+        self.result = result
+        self.env = jinja2.Environment(loader=jinja2.FileSystemLoader(template_dir))
+        
+        # Initialize tiktoken for context management
+        try:
+            self.encoding = tiktoken.get_encoding("cl100k_base") # GPT-4/3.5-turbo encoding
+        except Exception:
+            self.encoding = None
+            
+        # Initialize tree-sitter for precision extraction
+        try:
+            self.PY_LANGUAGE = Language(tree_sitter_python.language())
+            self.parser = Parser(self.PY_LANGUAGE)
+        except Exception:
+            self.parser = None
+        
+    def generate_prompts(self) -> Dict[str, str]:
+        """Generate a prompt for each detected code smell."""
+        prompts = {}
+        
+        for i, smell in enumerate(self.result.smells):
+            prompt = self._generate_prompt_for_smell(smell)
+            if prompt:
+                # Truncate prompt if it exceeds token limit (e.g., 4000 tokens)
+                if self.encoding:
+                    tokens = self.encoding.encode(prompt)
+                    if len(tokens) > 4000:
+                        prompt = self.encoding.decode(tokens[:3800]) + "\n\n... (prompt truncated due to length) ..."
+                
+                # Use a unique name for each prompt
+                filename = f"{i+1:02d}_{smell.type}_{smell.name.lower().replace(' ', '_').replace(':', '')}.md"
+                prompts[filename] = prompt
+        return prompts
+        
+    def _generate_prompt_for_smell(self, smell: CodeSmell) -> Optional[str]:
+        """Generate a single prompt from a CodeSmell."""
+        template_name = self._get_template_for_type(smell.type)
+        if not template_name:
+            return None
+            
+        try:
+            template = self.env.get_template(template_name)
+            context = self._build_context_for_smell(smell)
+            return template.render(**context)
+        except Exception as e:
+            print(f"Error generating prompt for {smell.name}: {e}")
+            return None
+            
+    def _get_template_for_type(self, smell_type: str) -> Optional[str]:
+        """Map smell type to Jinja2 template filename."""
+        mapping = {
+            "god_function": "extract_method.md",
+            "feature_envy": "move_method.md",
+            "data_clump": "move_method.md", 
+            "shotgun_surgery": "extract_method.md",
+            "bottleneck": "extract_method.md",
+            "circular_dependency": "move_method.md"
+        }
+        return mapping.get(smell_type)
+        
+    def _build_context_for_smell(self, smell: CodeSmell) -> Dict[str, Any]:
+        """Prepare context data for the Jinja2 template."""
+        # Extract source code for context
+        source_code = self._get_source_context(smell.file, smell.line)
+        
+        # Prepare metrics
+        metrics = self.result.metrics.get(smell.name.split(': ')[-1], {}) # Heuristic to find function name
+        if not metrics and 'function' in smell.context:
+            metrics = self.result.metrics.get(smell.context['function'], {})
+
+        # Prepare mutations
+        mutations = [m for m in self.result.mutations if m.scope in (smell.name.split(': ')[-1], smell.context.get('function'))]
+        mutations_summary = f"{len(mutations)} modifications recorded: {', '.join(set([m.variable for m in mutations[:5]]))}..."
+
+        context = {
+            "target_function": smell.name.split(': ')[-1],
+            "reason": smell.description,
+            "metrics": metrics,
+            "mutations_context": mutations_summary,
+            "source_file": smell.file,
+            "start_line": smell.line,
+            "end_line": smell.line + 20, # Heuristic: end of function or next 20 lines
+            "source_code": source_code,
+            "instruction": self._get_instruction_for_smell(smell),
+            # move_method specific
+            "source_module": smell.file.split('/')[-1].replace('.py', ''),
+            "target_module": smell.context.get('foreign_mutations', ["other_module"])[0].split('.')[0] if smell.type == "feature_envy" else "other_module",
+            "foreign_mutations": ", ".join(smell.context.get('foreign_mutations', [])),
+            "foreign_mutations_context": f"This code mutates state in {', '.join(set([v.split('.')[0] for v in smell.context.get('foreign_mutations', []) if '.' in v]))}",
+            "dependencies": ", ".join(set([m.variable for m in mutations if '.' in m.variable])),
+            "reachability": self.result.functions.get(smell.name.split(': ')[-1], {}).reachability if hasattr(self.result.functions.get(smell.name.split(': ')[-1]), 'reachability') else "unknown"
+        }
+        return context
+
+    def _get_source_context(self, file_path: str, start_line: int, max_lines: int = 50) -> str:
+        """Read source code lines from a file."""
+        if not os.path.exists(file_path):
+            return "# Source file not found."
+            
+        try:
+            with open(file_path, 'r') as f:
+                content = f.read()
+                
+            # If tree-sitter is available, use it to accurately find function boundaries
+            if self.parser and "method" not in file_path: # simplified check
+                tree = self.parser.parse(bytes(content, "utf8"))
+                root_node = tree.root_node
+                
+                # Simple function extraction using tree-sitter
+                # (Ideally we'd search for the function node at start_line)
+                lines = content.splitlines()
+                start = max(0, start_line - 1)
+                end = min(len(lines), start + max_lines)
+                return "\n".join(lines[start:end])
+            else:
+                lines = content.splitlines()
+                start = max(0, start_line - 1)
+                end = min(len(lines), start + max_lines)
+                return "\n".join(lines[start:end])
+        except Exception as e:
+            return f"# Error reading source: {e}"
+
+    def _get_instruction_for_smell(self, smell: CodeSmell) -> str:
+        """Generate specific instruction based on smell type."""
+        if smell.type == "god_function":
+            return f"Wyekstrahuj mniejsze, spójne metody z funkcji {smell.name.split(': ')[-1]}. Skup się na wydzieleniu operacji o największej liczbie mutacji."
+        elif smell.type == "feature_envy":
+            return f"Przenieś metodę {smell.name.split(': ')[-1]} do modułu, który posiada większość używanych w niej danych. Zmniejsz coupling między modułami."
+        elif smell.type == "bottleneck":
+            return f"Funkcja {smell.name.split(': ')[-1]} jest wąskim gardłem strukturalnym. Wyekstrahuj z niej niezależne części pomocnicze, aby ułatwić zrozumienie przepływu."
+        elif smell.type == "circular_dependency":
+            return f"Wykryto cykl zależności. Przenieś część logiki do nowego modułu lub użyj interfejsu, aby przerwać cykl."
+        return "Zrefaktoryzuj ten fragment kodu, aby poprawić jego strukturę i zmniejszyć złożoność."

code2flow/visualizers/graph.py

@@ -0,0 +1,196 @@
+"""Graph visualization using NetworkX and matplotlib."""
+
+import matplotlib.pyplot as plt
+import matplotlib.patches as patches
+import networkx as nx
+from typing import Dict
+from pathlib import Path
+
+from ..core.models import AnalysisResult
+from ..core.config import NODE_COLORS
+
+
+class GraphVisualizer:
+    """Visualize analysis results as graphs."""
+    
+    def __init__(self, result: AnalysisResult):
+        self.result = result
+        self.graph = nx.DiGraph()
+        self._build_graph()
+        
+    def _build_graph(self):
+        """Build NetworkX graph from analysis result."""
+        # Add nodes
+        for node_id, node in self.result.nodes.items():
+            color = NODE_COLORS.get(node.type, '#757575')
+            self.graph.add_node(
+                node_id,
+                label=node.label[:30],
+                type=node.type,
+                color=color,
+                function=node.function
+            )
+            
+        # Add edges
+        for edge in self.result.edges:
+            self.graph.add_edge(
+                edge.source,
+                edge.target,
+                edge_type=edge.edge_type,
+                conditions=edge.conditions
+            )
+            
+    def visualize_cfg(self, filepath: str, layout: str = 'spring'):
+        """Create control flow visualization."""
+        plt.figure(figsize=(16, 12))
+        
+        # Choose layout
+        if layout == 'spring':
+            pos = nx.spring_layout(self.graph, k=2, iterations=50, seed=42)
+        elif layout == 'hierarchical':
+            pos = self._hierarchical_layout()
+        elif layout == 'kamada':
+            pos = nx.kamada_kawai_layout(self.graph)
+        else:
+            pos = nx.shell_layout(self.graph)
+            
+        # Get node colors
+        node_colors = []
+        for node_id in self.graph.nodes():
+            node_type = self.graph.nodes[node_id].get('type', 'DEFAULT')
+            node_colors.append(NODE_COLORS.get(node_type, '#757575'))
+            
+        # Draw graph
+        nx.draw_networkx_nodes(
+            self.graph, pos,
+            node_color=node_colors,
+            node_size=600,
+            alpha=0.8,
+            edgecolors='white',
+            linewidths=2
+        )
+        
+        nx.draw_networkx_edges(
+            self.graph, pos,
+            alpha=0.4,
+            arrows=True,
+            arrowsize=15,
+            arrowstyle='->',
+            edge_color='#666666',
+            width=1.5
+        )
+        
+        # Draw labels for important nodes
+        labels = {}
+        for node_id in self.graph.nodes():
+            node_data = self.graph.nodes[node_id]
+            if node_data.get('type') in ['FUNC', 'IF', 'CALL']:
+                label = node_data.get('label', '') or ''
+                labels[node_id] = label[:25]
+                
+        nx.draw_networkx_labels(
+            self.graph, pos, labels,
+            font_size=8,
+            font_color='white',
+            font_weight='bold'
+        )
+        
+        # Add legend
+        legend_elements = [
+            patches.Patch(color=NODE_COLORS['FUNC'], label='Function'),
+            patches.Patch(color=NODE_COLORS['CALL'], label='Call'),
+            patches.Patch(color=NODE_COLORS['IF'], label='Decision'),
+            patches.Patch(color=NODE_COLORS['FOR'], label='Loop'),
+            patches.Patch(color=NODE_COLORS['RETURN'], label='Return')
+        ]
+        plt.legend(handles=legend_elements, loc='upper right', fontsize=10)
+        
+        plt.title('Control Flow Graph', fontsize=16, fontweight='bold')
+        plt.axis('off')
+        plt.tight_layout()
+        plt.savefig(filepath, dpi=300, bbox_inches='tight', facecolor='white')
+        plt.close()
+        
+    def visualize_call_graph(self, filepath: str):
+        """Visualize call graph."""
+        # Build call graph
+        call_graph = nx.DiGraph()
+        
+        for func_name, func_info in self.result.functions.items():
+            short_name = func_name.split('.')[-1]
+            call_graph.add_node(func_name, label=short_name)
+            
+            for callee in func_info.calls:
+                call_graph.add_edge(func_name, callee)
+                
+        if len(call_graph.nodes()) == 0:
+            return  # Nothing to visualize
+            
+        plt.figure(figsize=(14, 10))
+        
+        # Layout
+        pos = nx.spring_layout(call_graph, k=1.5, iterations=50, seed=42)
+        
+        # Node sizes based on calls
+        node_sizes = []
+        for node in call_graph.nodes():
+            out_degree = call_graph.out_degree(node)
+            in_degree = call_graph.in_degree(node)
+            size = 300 + (out_degree + in_degree) * 150
+            node_sizes.append(size)
+            
+        # Draw
+        nx.draw_networkx_nodes(
+            call_graph, pos,
+            node_color='#4CAF50',
+            node_size=node_sizes,
+            alpha=0.8,
+            edgecolors='white',
+            linewidths=2
+        )
+        
+        nx.draw_networkx_edges(
+            call_graph, pos,
+            alpha=0.5,
+            arrows=True,
+            arrowsize=20,
+            edge_color='#2196F3',
+            width=2
+        )
+        
+        # Labels
+        labels = {n: (call_graph.nodes[n].get('label') or n)[:20] for n in call_graph.nodes()}
+        nx.draw_networkx_labels(
+            call_graph, pos, labels,
+            font_size=9,
+            font_weight='bold'
+        )
+        
+        plt.title('Function Call Graph', fontsize=16, fontweight='bold')
+        plt.axis('off')
+        plt.tight_layout()
+        plt.savefig(filepath, dpi=300, bbox_inches='tight', facecolor='white')
+        plt.close()
+        
+    def _hierarchical_layout(self) -> Dict:
+        """Create hierarchical layout grouped by function."""
+        from collections import defaultdict
+        
+        # Group nodes by function
+        function_groups = defaultdict(list)
+        for node_id, node in self.result.nodes.items():
+            func = node.function or '__global__'
+            function_groups[func].append(node_id)
+            
+        # Position nodes
+        pos = {}
+        y_offset = 0
+        
+        for func_name, nodes in sorted(function_groups.items()):
+            for i, node_id in enumerate(nodes):
+                x = i * 2
+                y = -y_offset
+                pos[node_id] = (x, y)
+            y_offset += 3
+            
+        return pos