code_interpreter.py

import ast
from typing import Dict, List, Any, Optional
import re

class CodeInterpreter:
    def __init__(self):
        self.language = 'python'  # Only support Python
        
    def analyze_code(self, code: str) -> Dict[str, Any]:
        """
        Analyze Python code and extract key information about its structure and functionality.
        """
        try:
            return self._analyze_python_code(code)
        except Exception as e:
            return {"error": f"Python code analysis failed: {str(e)}"}
    
    def _analyze_python_code(self, code: str) -> Dict[str, Any]:
        """
        Analyze Python code using AST.
        """
        try:
            tree = ast.parse(code)
            analysis = {
                "imports": [],
                "functions": [],
                "classes": [],
                "variables": [],
                "complexity": 0,
                "docstrings": [],
                "decorators": []
            }
            
            for node in ast.walk(tree):
                if isinstance(node, ast.Import):
                    for name in node.names:
                        analysis["imports"].append(name.name)
                elif isinstance(node, ast.ImportFrom):
                    analysis["imports"].append(f"{node.module}.{node.names[0].name}")
                elif isinstance(node, ast.FunctionDef):
                    func_info = {
                        "name": node.name,
                        "args": [arg.arg for arg in node.args.args],
                        "returns": self._get_return_type(node),
                        "complexity": self._calculate_complexity(node),
                        "docstring": ast.get_docstring(node),
                        "decorators": [d.id for d in node.decorator_list if isinstance(d, ast.Name)]
                    }
                    analysis["functions"].append(func_info)
                elif isinstance(node, ast.ClassDef):
                    class_info = {
                        "name": node.name,
                        "methods": [],
                        "bases": [base.id for base in node.bases if isinstance(base, ast.Name)],
                        "docstring": ast.get_docstring(node),
                        "decorators": [d.id for d in node.decorator_list if isinstance(d, ast.Name)]
                    }
                    for item in node.body:
                        if isinstance(item, ast.FunctionDef):
                            class_info["methods"].append(item.name)
                    analysis["classes"].append(class_info)
                elif isinstance(node, ast.Assign):
                    for target in node.targets:
                        if isinstance(target, ast.Name):
                            analysis["variables"].append(target.id)
                elif isinstance(node, ast.Expr) and isinstance(node.value, ast.Str):
                    analysis["docstrings"].append(node.value.s)
            
            analysis["complexity"] = sum(func["complexity"] for func in analysis["functions"])
            return analysis
            
        except Exception as e:
            return {"error": f"Python code analysis failed: {str(e)}"}
    
    def _get_return_type(self, node: ast.FunctionDef) -> Optional[str]:
        """Extract return type annotation if present."""
        if node.returns:
            if isinstance(node.returns, ast.Name):
                return node.returns.id
            elif isinstance(node.returns, ast.Subscript):
                return f"{node.returns.value.id}[{node.returns.slice.value.id}]"
        return None
    
    def _calculate_complexity(self, node: ast.AST) -> int:
        """Calculate cyclomatic complexity of a function."""
        complexity = 1
        for child in ast.walk(node):
            if isinstance(child, (ast.If, ast.While, ast.For, ast.Try, ast.ExceptHandler)):
                complexity += 1
        return complexity
    
    def suggest_improvements(self, analysis: Dict[str, Any]) -> List[str]:
        """
        Suggest code improvements based on analysis.
        """
        suggestions = []
        
        # Check function complexity
        for func in analysis.get("functions", []):
            if func["complexity"] > 10:
                suggestions.append(f"Function '{func['name']}' is too complex (complexity: {func['complexity']}). Consider breaking it down into smaller functions.")
        
        # Check for missing type hints
        for func in analysis.get("functions", []):
            if not func["returns"]:
                suggestions.append(f"Function '{func['name']}' is missing return type annotation.")
        
        # Check for missing docstrings
        for func in analysis.get("functions", []):
            if not func["docstring"]:
                suggestions.append(f"Function '{func['name']}' is missing a docstring.")
        
        # Check for unused imports
        if len(analysis.get("imports", [])) > 10:
            suggestions.append("Consider removing unused imports to improve code clarity.")
        
        # Check for long functions
        for func in analysis.get("functions", []):
            if len(func["args"]) > 5:
                suggestions.append(f"Function '{func['name']}' has too many parameters ({len(func['args'])}). Consider using a data class or dictionary.")
        
        return suggestions
    
    def extract_code_context(self, code: str, line_number: int) -> Dict[str, Any]:
        """
        Extract context around a specific line of code.
        """
        lines = code.split('\n')
        context = {
            "line": lines[line_number - 1] if 0 <= line_number - 1 < len(lines) else "",
            "before": lines[max(0, line_number - 3):line_number - 1],
            "after": lines[line_number:min(len(lines), line_number + 3)],
            "indentation": len(re.match(r'^\s*', lines[line_number - 1]).group()) if 0 <= line_number - 1 < len(lines) else 0
        }
        return context