codeshift 0.2.0__py3-none-any.whl

codeshift/utils/llm_client.py

@@ -0,0 +1,221 @@
+"""Anthropic Claude client wrapper for LLM-based migrations."""
+
+import os
+from dataclasses import dataclass
+
+from anthropic import Anthropic
+
+
+@dataclass
+class LLMResponse:
+    """Response from the LLM."""
+
+    content: str
+    model: str
+    usage: dict
+    success: bool
+    error: str | None = None
+
+
+class LLMClient:
+    """Client for interacting with Anthropic's Claude API."""
+
+    DEFAULT_MODEL = "claude-sonnet-4-20250514"
+    MAX_TOKENS = 4096
+
+    def __init__(
+        self,
+        api_key: str | None = None,
+        model: str | None = None,
+    ):
+        """Initialize the LLM client.
+
+        Args:
+            api_key: Anthropic API key. Defaults to ANTHROPIC_API_KEY env var.
+            model: Model to use. Defaults to claude-sonnet-4-20250514.
+        """
+        self.api_key = api_key or os.environ.get("ANTHROPIC_API_KEY")
+        self.model = model or self.DEFAULT_MODEL
+        self._client: Anthropic | None = None
+
+    @property
+    def client(self) -> Anthropic:
+        """Get or create the Anthropic client."""
+        if self._client is None:
+            if not self.api_key:
+                raise ValueError(
+                    "Anthropic API key not found. Set ANTHROPIC_API_KEY environment variable "
+                    "or pass api_key to LLMClient."
+                )
+            self._client = Anthropic(api_key=self.api_key)
+        return self._client
+
+    @property
+    def is_available(self) -> bool:
+        """Check if the LLM client is available (API key is set)."""
+        return bool(self.api_key)
+
+    def generate(
+        self,
+        prompt: str,
+        system_prompt: str | None = None,
+        max_tokens: int | None = None,
+        temperature: float = 0.0,
+    ) -> LLMResponse:
+        """Generate a response from the LLM.
+
+        Args:
+            prompt: The user prompt
+            system_prompt: Optional system prompt
+            max_tokens: Maximum tokens in response
+            temperature: Sampling temperature (0.0 for deterministic)
+
+        Returns:
+            LLMResponse with the generated content
+        """
+        if not self.is_available:
+            return LLMResponse(
+                content="",
+                model=self.model,
+                usage={},
+                success=False,
+                error="API key not configured",
+            )
+
+        try:
+            response = self.client.messages.create(
+                model=self.model,
+                max_tokens=max_tokens or self.MAX_TOKENS,
+                temperature=temperature,
+                system=system_prompt or "",
+                messages=[{"role": "user", "content": prompt}],
+            )
+
+            content = ""
+            for block in response.content:
+                if hasattr(block, "text"):
+                    content += block.text
+
+            return LLMResponse(
+                content=content,
+                model=response.model,
+                usage={
+                    "input_tokens": response.usage.input_tokens,
+                    "output_tokens": response.usage.output_tokens,
+                },
+                success=True,
+            )
+
+        except Exception as e:
+            return LLMResponse(
+                content="",
+                model=self.model,
+                usage={},
+                success=False,
+                error=str(e),
+            )
+
+    def migrate_code(
+        self,
+        code: str,
+        library: str,
+        from_version: str,
+        to_version: str,
+        context: str | None = None,
+    ) -> LLMResponse:
+        """Use the LLM to migrate code.
+
+        Args:
+            code: The source code to migrate
+            library: The library being upgraded
+            from_version: Current version
+            to_version: Target version
+            context: Optional context about the migration
+
+        Returns:
+            LLMResponse with the migrated code
+        """
+        system_prompt = f"""You are an expert Python developer specializing in code migrations.
+Your task is to migrate Python code from {library} v{from_version} to v{to_version}.
+
+Guidelines:
+1. Only modify code that needs to change for the migration
+2. Preserve all comments, formatting, and code style where possible
+3. Add brief inline comments explaining non-obvious changes
+4. If you're unsure about a change, add a TODO comment
+5. Return ONLY the migrated code, no explanations before or after
+
+Important {library} v{from_version} to v{to_version} changes:
+- Config class -> model_config = ConfigDict(...)
+- @validator -> @field_validator with @classmethod
+- @root_validator -> @model_validator with @classmethod
+- .dict() -> .model_dump()
+- .json() -> .model_dump_json()
+- .schema() -> .model_json_schema()
+- .parse_obj() -> .model_validate()
+- .parse_raw() -> .model_validate_json()
+- .copy() -> .model_copy()
+- orm_mode -> from_attributes
+- Field(regex=...) -> Field(pattern=...)
+"""
+
+        prompt = f"""Migrate the following Python code from {library} v{from_version} to v{to_version}.
+
+{f"Context: {context}" if context else ""}
+
+Code to migrate:
+```python
+{code}
+```
+
+Return only the migrated Python code:"""
+
+        return self.generate(prompt, system_prompt=system_prompt)
+
+    def explain_change(
+        self,
+        original: str,
+        transformed: str,
+        library: str,
+    ) -> LLMResponse:
+        """Use the LLM to explain a migration change.
+
+        Args:
+            original: Original code
+            transformed: Transformed code
+            library: The library being upgraded
+
+        Returns:
+            LLMResponse with the explanation
+        """
+        system_prompt = """You are an expert Python developer.
+Explain code changes clearly and concisely for other developers.
+Focus on the 'why' not just the 'what'."""
+
+        prompt = f"""Explain the following {library} migration change:
+
+Original:
+```python
+{original}
+```
+
+Migrated:
+```python
+{transformed}
+```
+
+Provide a brief explanation (2-3 sentences) of what changed and why:"""
+
+        return self.generate(prompt, system_prompt=system_prompt, max_tokens=500)
+
+
+# Singleton instance for convenience
+_default_client: LLMClient | None = None
+
+
+def get_llm_client() -> LLMClient:
+    """Get the default LLM client instance."""
+    global _default_client
+    if _default_client is None:
+        _default_client = LLMClient()
+    return _default_client

codeshift/validator/__init__.py

@@ -0,0 +1,6 @@
+"""Validator module for checking transformed code."""
+
+from codeshift.validator.syntax_checker import SyntaxChecker, SyntaxCheckResult
+from codeshift.validator.test_runner import TestResult, TestRunner
+
+__all__ = ["SyntaxChecker", "SyntaxCheckResult", "TestRunner", "TestResult"]

codeshift/validator/syntax_checker.py

@@ -0,0 +1,183 @@
+"""Syntax checker for validating transformed code."""
+
+import ast
+import sys
+from dataclasses import dataclass, field
+from pathlib import Path
+
+
+@dataclass
+class SyntaxIssue:
+    """Represents a syntax error in code."""
+
+    message: str
+    line_number: int
+    column: int
+    line_text: str | None = None
+
+
+@dataclass
+class SyntaxCheckResult:
+    """Result of a syntax check."""
+
+    is_valid: bool
+    file_path: Path | None = None
+    errors: list[SyntaxIssue] = field(default_factory=list)
+    warnings: list[str] = field(default_factory=list)
+
+    @property
+    def error_count(self) -> int:
+        """Get the number of errors."""
+        return len(self.errors)
+
+
+class SyntaxChecker:
+    """Validates Python code syntax."""
+
+    def __init__(self, python_version: tuple[int, int] | None = None):
+        """Initialize the syntax checker.
+
+        Args:
+            python_version: Target Python version as (major, minor).
+                          Defaults to current Python version.
+        """
+        if python_version is None:
+            python_version = (sys.version_info.major, sys.version_info.minor)
+        self.python_version = python_version
+
+    def check_code(self, source_code: str, filename: str = "<string>") -> SyntaxCheckResult:
+        """Check if source code has valid Python syntax.
+
+        Args:
+            source_code: The Python source code to check
+            filename: Optional filename for error messages
+
+        Returns:
+            SyntaxCheckResult with validation status
+        """
+        try:
+            # First, try to compile the code
+            compile(source_code, filename, "exec")
+
+            # Then parse with AST for more detailed checking
+            ast.parse(source_code, filename=filename)
+
+            return SyntaxCheckResult(is_valid=True)
+
+        except SyntaxError as e:
+            error = SyntaxIssue(
+                message=str(e.msg) if hasattr(e, "msg") else str(e),
+                line_number=e.lineno or 0,
+                column=e.offset or 0,
+                line_text=e.text,
+            )
+            return SyntaxCheckResult(
+                is_valid=False,
+                errors=[error],
+            )
+
+    def check_file(self, file_path: Path) -> SyntaxCheckResult:
+        """Check if a Python file has valid syntax.
+
+        Args:
+            file_path: Path to the Python file
+
+        Returns:
+            SyntaxCheckResult with validation status
+        """
+        try:
+            source_code = file_path.read_text()
+        except Exception as e:
+            return SyntaxCheckResult(
+                is_valid=False,
+                file_path=file_path,
+                errors=[
+                    SyntaxIssue(
+                        message=f"Could not read file: {e}",
+                        line_number=0,
+                        column=0,
+                    )
+                ],
+            )
+
+        result = self.check_code(source_code, str(file_path))
+        result.file_path = file_path
+        return result
+
+    def check_directory(
+        self, directory: Path, exclude_patterns: list[str] | None = None
+    ) -> list[SyntaxCheckResult]:
+        """Check all Python files in a directory.
+
+        Args:
+            directory: Path to the directory
+            exclude_patterns: Glob patterns to exclude
+
+        Returns:
+            List of SyntaxCheckResult for each file with errors
+        """
+        import fnmatch
+
+        exclude_patterns = exclude_patterns or []
+        results = []
+
+        for file_path in directory.rglob("*.py"):
+            relative_path = str(file_path.relative_to(directory))
+
+            # Check exclude patterns
+            excluded = False
+            for pattern in exclude_patterns:
+                if fnmatch.fnmatch(relative_path, pattern):
+                    excluded = True
+                    break
+
+            if excluded:
+                continue
+
+            result = self.check_file(file_path)
+            if not result.is_valid:
+                results.append(result)
+
+        return results
+
+    def validate_transform(self, original: str, transformed: str) -> tuple[bool, list[str]]:
+        """Validate that a transformation didn't break syntax.
+
+        Args:
+            original: Original source code
+            transformed: Transformed source code
+
+        Returns:
+            Tuple of (is_valid, list of issues)
+        """
+        issues = []
+
+        # Check original syntax (should be valid)
+        original_result = self.check_code(original, "<original>")
+        if not original_result.is_valid:
+            issues.append("Original code has syntax errors")
+
+        # Check transformed syntax
+        transformed_result = self.check_code(transformed, "<transformed>")
+        if not transformed_result.is_valid:
+            for error in transformed_result.errors:
+                issues.append(f"Line {error.line_number}: {error.message}")
+            return False, issues
+
+        return True, issues
+
+
+def quick_syntax_check(source_code: str) -> bool:
+    """Quick check if code has valid Python syntax.
+
+    Args:
+        source_code: The Python source code to check
+
+    Returns:
+        True if syntax is valid, False otherwise
+    """
+    try:
+        compile(source_code, "<string>", "exec")
+        return True
+    except SyntaxError:
+        return False

codeshift/validator/test_runner.py

@@ -0,0 +1,224 @@
+"""Test runner for validating migrations."""
+
+import subprocess
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+
+
+@dataclass
+class TestResult:
+    """Result of running tests."""
+
+    success: bool
+    exit_code: int
+    stdout: str = ""
+    stderr: str = ""
+    tests_run: int = 0
+    tests_passed: int = 0
+    tests_failed: int = 0
+    tests_skipped: int = 0
+    duration: float = 0.0
+    error_message: str | None = None
+
+    @property
+    def summary(self) -> str:
+        """Get a summary string of the test results."""
+        if self.success:
+            return f"✓ {self.tests_passed}/{self.tests_run} tests passed"
+        return f"✗ {self.tests_failed}/{self.tests_run} tests failed"
+
+
+class TestRunner:
+    """Runs project tests to validate migrations."""
+
+    def __init__(
+        self,
+        project_path: Path,
+        test_command: list[str] | None = None,
+        timeout: int = 300,
+    ):
+        """Initialize the test runner.
+
+        Args:
+            project_path: Path to the project root
+            test_command: Custom test command. Defaults to pytest detection.
+            timeout: Maximum time in seconds to run tests
+        """
+        self.project_path = project_path
+        self.test_command = test_command or self._detect_test_command()
+        self.timeout = timeout
+
+    def _detect_test_command(self) -> list[str]:
+        """Detect the appropriate test command for the project."""
+        # Check for pytest
+        if (
+            (self.project_path / "pytest.ini").exists()
+            or (self.project_path / "pyproject.toml").exists()
+            or (self.project_path / "tests").exists()
+        ):
+            return [sys.executable, "-m", "pytest", "-v", "--tb=short"]
+
+        # Check for unittest
+        if (self.project_path / "tests").exists():
+            return [sys.executable, "-m", "unittest", "discover", "-v"]
+
+        # Default to pytest
+        return [sys.executable, "-m", "pytest", "-v", "--tb=short"]
+
+    def run(
+        self,
+        specific_tests: list[str] | None = None,
+        extra_args: list[str] | None = None,
+    ) -> TestResult:
+        """Run the project tests.
+
+        Args:
+            specific_tests: List of specific test files or patterns to run
+            extra_args: Additional arguments to pass to the test runner
+
+        Returns:
+            TestResult with the outcome
+        """
+        command = self.test_command.copy()
+
+        if extra_args:
+            command.extend(extra_args)
+
+        if specific_tests:
+            command.extend(specific_tests)
+
+        try:
+            result = subprocess.run(
+                command,
+                cwd=self.project_path,
+                capture_output=True,
+                text=True,
+                timeout=self.timeout,
+            )
+
+            # Parse pytest output
+            tests_run, tests_passed, tests_failed, tests_skipped = self._parse_pytest_output(
+                result.stdout + result.stderr
+            )
+
+            return TestResult(
+                success=result.returncode == 0,
+                exit_code=result.returncode,
+                stdout=result.stdout,
+                stderr=result.stderr,
+                tests_run=tests_run,
+                tests_passed=tests_passed,
+                tests_failed=tests_failed,
+                tests_skipped=tests_skipped,
+            )
+
+        except subprocess.TimeoutExpired:
+            return TestResult(
+                success=False,
+                exit_code=-1,
+                error_message=f"Tests timed out after {self.timeout} seconds",
+            )
+        except FileNotFoundError as e:
+            return TestResult(
+                success=False,
+                exit_code=-1,
+                error_message=f"Test command not found: {e}",
+            )
+        except Exception as e:
+            return TestResult(
+                success=False,
+                exit_code=-1,
+                error_message=f"Error running tests: {e}",
+            )
+
+    def _parse_pytest_output(self, output: str) -> tuple[int, int, int, int]:
+        """Parse pytest output to extract test counts.
+
+        Args:
+            output: Combined stdout and stderr from pytest
+
+        Returns:
+            Tuple of (total, passed, failed, skipped)
+        """
+        import re
+
+        # Look for pytest summary line like "5 passed, 2 failed, 1 skipped"
+        # or "1 passed in 0.05s"
+        passed = 0
+        failed = 0
+        skipped = 0
+        errors = 0
+
+        # Match patterns like "5 passed", "2 failed", etc.
+        passed_match = re.search(r"(\d+) passed", output)
+        if passed_match:
+            passed = int(passed_match.group(1))
+
+        failed_match = re.search(r"(\d+) failed", output)
+        if failed_match:
+            failed = int(failed_match.group(1))
+
+        skipped_match = re.search(r"(\d+) skipped", output)
+        if skipped_match:
+            skipped = int(skipped_match.group(1))
+
+        error_match = re.search(r"(\d+) error", output)
+        if error_match:
+            errors = int(error_match.group(1))
+
+        total = passed + failed + skipped + errors
+        return total, passed, failed + errors, skipped
+
+    def run_quick_check(self) -> TestResult:
+        """Run a quick smoke test (collection only, no execution).
+
+        Returns:
+            TestResult indicating if tests can be collected
+        """
+        command = self.test_command.copy()
+        command.extend(["--collect-only", "-q"])
+
+        try:
+            result = subprocess.run(
+                command,
+                cwd=self.project_path,
+                capture_output=True,
+                text=True,
+                timeout=60,
+            )
+
+            # Count collected tests
+            tests_collected = 0
+            for line in result.stdout.splitlines():
+                if "test" in line.lower() and "::" in line:
+                    tests_collected += 1
+
+            return TestResult(
+                success=result.returncode == 0,
+                exit_code=result.returncode,
+                stdout=result.stdout,
+                stderr=result.stderr,
+                tests_run=tests_collected,
+            )
+
+        except Exception as e:
+            return TestResult(
+                success=False,
+                exit_code=-1,
+                error_message=f"Error collecting tests: {e}",
+            )
+
+
+def run_tests(project_path: Path, timeout: int = 300) -> TestResult:
+    """Convenience function to run tests for a project.
+
+    Args:
+        project_path: Path to the project
+        timeout: Maximum time in seconds
+
+    Returns:
+        TestResult with the outcome
+    """
+    runner = TestRunner(project_path, timeout=timeout)
+    return runner.run()