Build an AI Legacy Code Refactoring Tool with Python

mkdir ai-refactor && cd ai-refactor
pip install anthropic pathspec
export ANTHROPIC_API_KEY="sk-your-ezai-key"

import ast, re
from pathlib import Path
from dataclasses import dataclass, field

@dataclass
class CodeSmell:
    file: Path
    line: int
    category: str   # "sql_injection" | "no_types" | "god_function" | "bare_except"
    description: str
    severity: int    # 1-5, where 5 = critical

class CodeAnalyzer:
    SMELL_PATTERNS = [
        (r'"\s*SELECT.*\+\s*', "sql_injection", "String-concatenated SQL query", 5),
        (r'except\s*:', "bare_except", "Bare except catches everything", 3),
        (r'import \*', "wildcard_import", "Wildcard import pollutes namespace", 2),
        (r'\.format\(.*input\(', "unsafe_input", "User input in string format", 4),
    ]

    def analyze_file(self, path: Path) -> list[CodeSmell]:
        smells: list[CodeSmell] = []
        source = path.read_text()
        lines = source.splitlines()

        # Pattern-based detection
        for i, line in enumerate(lines, 1):
            for pattern, cat, desc, sev in self.SMELL_PATTERNS:
                if re.search(pattern, line):
                    smells.append(CodeSmell(path, i, cat, desc, sev))

        # AST-based detection: god functions (>50 lines)
        try:
            tree = ast.parse(source)
            for node in ast.walk(tree):
                if isinstance(node, ast.FunctionDef):
                    length = node.end_lineno - node.lineno
                    if length > 50:
                        smells.append(CodeSmell(
                            path, node.lineno, "god_function",
                            f"{node.name}() is {length} lines", 3
                        ))
                    # Missing type hints
                    if not node.returns:
                        smells.append(CodeSmell(
                            path, node.lineno, "no_types",
                            f"{node.name}() missing return type", 2
                        ))
        except SyntaxError:
            pass

        return smells

import anthropic

client = anthropic.Anthropic(
    base_url="https://ezaiapi.com"
)

REFACTOR_PROMPT = """You are a senior Python engineer refactoring legacy code.

DETECTED ISSUES:
{smells}

RULES:
1. Fix ONLY the listed issues — don't rewrite unrelated code
2. Preserve all existing behavior and function signatures
3. Add type hints to every function (params + return)
4. Replace string-concatenated SQL with parameterized queries
5. Replace bare except with specific exception types
6. Break god functions into smaller, focused functions
7. Add docstrings to public functions
8. Use modern Python (dataclasses, f-strings, pathlib)

Return ONLY the refactored Python code. No explanations."""

def refactor_file(source: str, smells: list[CodeSmell]) -> str:
    smell_desc = "\n".join(
        f"- Line {s.line}: [{s.category}] {s.description}"
        for s in smells
    )

    response = client.messages.create(
        model="claude-sonnet-4-5",
        max_tokens=8192,
        messages=[
            {"role": "user", "content": f"""{REFACTOR_PROMPT.format(smells=smell_desc)}

SOURCE CODE:
```python
{source}
```"""}
        ]
    )

    # Extract code from markdown fence
    text = response.content[0].text
    if "```python" in text:
        text = text.split("```python")[1].split("```")[0]
    return text.strip()

import subprocess

def validate_refactor(original: str, refactored: str, path: Path) -> bool:
    # 1. Check it parses
    try:
        ast.parse(refactored)
    except SyntaxError as e:
        print(f"❌ Syntax error in refactored code: {e}")
        return False

    # 2. Verify function signatures preserved
    orig_funcs = {
        n.name: len(n.args.args)
        for n in ast.walk(ast.parse(original))
        if isinstance(n, ast.FunctionDef)
    }
    new_funcs = {
        n.name: len(n.args.args)
        for n in ast.walk(ast.parse(refactored))
        if isinstance(n, ast.FunctionDef)
    }
    for name, argc in orig_funcs.items():
        if name not in new_funcs:
            print(f"❌ Function {name}() was removed")
            return False
        if new_funcs[name] != argc:
            print(f"⚠️  {name}() args changed: {argc} → {new_funcs[name]}")

    # 3. Run existing tests if they exist
    test_file = path.parent / f"test_{path.name}"
    if test_file.exists():
        result = subprocess.run(
            ["python", "-m", "pytest", str(test_file), "-q"],
            capture_output=True, text=True
        )
        if result.returncode != 0:
            print(f"❌ Tests failed:\n{result.stdout}")
            return False

    return True

import argparse, json, time

def main():
    parser = argparse.ArgumentParser(description="AI-powered legacy code refactoring")
    parser.add_argument("path", type=Path, help="Directory or file to refactor")
    parser.add_argument("--dry-run", action="store_true", help="Analyze only, don't apply changes")
    parser.add_argument("--min-severity", type=int, default=2, help="Minimum smell severity (1-5)")
    args = parser.parse_args()

    analyzer = CodeAnalyzer()
    files = list(args.path.rglob("*.py")) if args.path.is_dir() else [args.path]
    results = {"refactored": 0, "skipped": 0, "failed": 0, "smells_fixed": 0}

    for f in files:
        smells = [s for s in analyzer.analyze_file(f)
                  if s.severity >= args.min_severity]
        if not smells:
            continue

        print(f"\n📄 {f} — {len(smells)} issues found")
        for s in smells:
            print(f"   L{s.line}: [{s.severity}] {s.description}")

        if args.dry_run:
            results["skipped"] += 1
            continue

        source = f.read_text()
        start = time.time()
        refactored = refactor_file(source, smells)
        elapsed = time.time() - start

        if validate_refactor(source, refactored, f):
            f.write_text(refactored)
            results["refactored"] += 1
            results["smells_fixed"] += len(smells)
            print(f"   ✅ Refactored in {elapsed:.1f}s")
        else:
            results["failed"] += 1
            print(f"   ❌ Validation failed — skipping")

    print(f"\n📊 Done: {results}")

if __name__ == "__main__":
    main()

# Analyze only — see what would change
python refactor.py ./legacy_app --dry-run

# Refactor everything with severity >= 3
python refactor.py ./legacy_app --min-severity 3

# Refactor a single file
python refactor.py ./legacy_app/db.py