codeshift 0.2.0__py3-none-any.whl

codeshift/migrator/ast_transforms.py

@@ -0,0 +1,256 @@
+"""Base transformer infrastructure for AST-based code migrations."""
+
+from dataclasses import dataclass, field
+from enum import Enum
+from pathlib import Path
+
+import libcst as cst
+
+
+class TransformStatus(Enum):
+    """Status of a transformation."""
+
+    SUCCESS = "success"
+    PARTIAL = "partial"  # Some transforms applied, some failed
+    FAILED = "failed"
+    NO_CHANGES = "no_changes"
+
+
+@dataclass
+class TransformChange:
+    """Represents a single code change made by a transform."""
+
+    description: str
+    line_number: int
+    original: str
+    replacement: str
+    transform_name: str
+    confidence: float = 1.0  # 0.0 to 1.0
+    notes: str | None = None
+
+
+@dataclass
+class TransformResult:
+    """Result of applying transforms to a file."""
+
+    file_path: Path
+    status: TransformStatus
+    original_code: str
+    transformed_code: str
+    changes: list[TransformChange] = field(default_factory=list)
+    errors: list[str] = field(default_factory=list)
+
+    @property
+    def has_changes(self) -> bool:
+        """Check if any changes were made."""
+        return self.original_code != self.transformed_code
+
+    @property
+    def change_count(self) -> int:
+        """Get the number of changes made."""
+        return len(self.changes)
+
+    def get_diff_lines(self) -> list[str]:
+        """Get a simple diff representation."""
+        import difflib
+
+        original_lines = self.original_code.splitlines(keepends=True)
+        transformed_lines = self.transformed_code.splitlines(keepends=True)
+
+        diff = difflib.unified_diff(
+            original_lines,
+            transformed_lines,
+            fromfile=f"{self.file_path} (original)",
+            tofile=f"{self.file_path} (transformed)",
+        )
+        return list(diff)
+
+
+class BaseTransformer(cst.CSTTransformer):
+    """Base class for LibCST transformers with change tracking."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.changes: list[TransformChange] = []
+        self.errors: list[str] = []
+        self._source_lines: list[str] = []
+
+    def set_source(self, source: str) -> None:
+        """Set the source code for reference during transforms."""
+        self._source_lines = source.splitlines()
+
+    def record_change(
+        self,
+        description: str,
+        line_number: int,
+        original: str,
+        replacement: str,
+        transform_name: str,
+        confidence: float = 1.0,
+        notes: str | None = None,
+    ) -> None:
+        """Record a change made by the transformer."""
+        self.changes.append(
+            TransformChange(
+                description=description,
+                line_number=line_number,
+                original=original,
+                replacement=replacement,
+                transform_name=transform_name,
+                confidence=confidence,
+                notes=notes,
+            )
+        )
+
+    def record_error(self, error: str) -> None:
+        """Record an error that occurred during transformation."""
+        self.errors.append(error)
+
+    def get_line(self, line_number: int) -> str:
+        """Get a specific line from the source."""
+        if 0 < line_number <= len(self._source_lines):
+            return self._source_lines[line_number - 1]
+        return ""
+
+
+def transform_file(
+    file_path: Path,
+    transformer: BaseTransformer,
+) -> TransformResult:
+    """Transform a file using the given transformer.
+
+    Args:
+        file_path: Path to the file to transform
+        transformer: The transformer to use
+
+    Returns:
+        TransformResult with the original and transformed code
+    """
+    try:
+        original_code = file_path.read_text()
+    except Exception as e:
+        return TransformResult(
+            file_path=file_path,
+            status=TransformStatus.FAILED,
+            original_code="",
+            transformed_code="",
+            errors=[f"Failed to read file: {e}"],
+        )
+
+    return transform_code(original_code, file_path, transformer)
+
+
+def transform_code(
+    source_code: str,
+    file_path: Path,
+    transformer: BaseTransformer,
+) -> TransformResult:
+    """Transform source code using the given transformer.
+
+    Args:
+        source_code: The source code to transform
+        file_path: Path for reference in results
+        transformer: The transformer to use
+
+    Returns:
+        TransformResult with the original and transformed code
+    """
+    try:
+        tree = cst.parse_module(source_code)
+    except cst.ParserSyntaxError as e:
+        return TransformResult(
+            file_path=file_path,
+            status=TransformStatus.FAILED,
+            original_code=source_code,
+            transformed_code=source_code,
+            errors=[f"Failed to parse file: {e}"],
+        )
+
+    transformer.set_source(source_code)
+
+    try:
+        transformed_tree = tree.visit(transformer)
+        transformed_code = transformed_tree.code
+    except Exception as e:
+        return TransformResult(
+            file_path=file_path,
+            status=TransformStatus.FAILED,
+            original_code=source_code,
+            transformed_code=source_code,
+            errors=[f"Transform failed: {e}"],
+        )
+
+    # Determine status
+    if transformer.errors:
+        if transformer.changes:
+            status = TransformStatus.PARTIAL
+        else:
+            status = TransformStatus.FAILED
+    elif transformer.changes:
+        status = TransformStatus.SUCCESS
+    else:
+        status = TransformStatus.NO_CHANGES
+
+    return TransformResult(
+        file_path=file_path,
+        status=status,
+        original_code=source_code,
+        transformed_code=transformed_code,
+        changes=transformer.changes,
+        errors=transformer.errors,
+    )
+
+
+def apply_transforms(
+    file_path: Path,
+    transformers: list[BaseTransformer],
+) -> TransformResult:
+    """Apply multiple transformers to a file in sequence.
+
+    Args:
+        file_path: Path to the file to transform
+        transformers: List of transformers to apply in order
+
+    Returns:
+        Combined TransformResult
+    """
+    try:
+        current_code = file_path.read_text()
+    except Exception as e:
+        return TransformResult(
+            file_path=file_path,
+            status=TransformStatus.FAILED,
+            original_code="",
+            transformed_code="",
+            errors=[f"Failed to read file: {e}"],
+        )
+
+    original_code = current_code
+    all_changes = []
+    all_errors = []
+
+    for transformer in transformers:
+        result = transform_code(current_code, file_path, transformer)
+        current_code = result.transformed_code
+        all_changes.extend(result.changes)
+        all_errors.extend(result.errors)
+
+    # Determine final status
+    if all_errors:
+        if all_changes:
+            status = TransformStatus.PARTIAL
+        else:
+            status = TransformStatus.FAILED
+    elif all_changes:
+        status = TransformStatus.SUCCESS
+    else:
+        status = TransformStatus.NO_CHANGES
+
+    return TransformResult(
+        file_path=file_path,
+        status=status,
+        original_code=original_code,
+        transformed_code=current_code,
+        changes=all_changes,
+        errors=all_errors,
+    )

codeshift/migrator/engine.py

@@ -0,0 +1,395 @@
+"""Migration engine with tiered approach."""
+
+from collections.abc import Callable
+from pathlib import Path
+
+from codeshift.knowledge import (
+    Confidence,
+    GeneratedKnowledgeBase,
+    is_tier_1_library,
+)
+from codeshift.migrator.ast_transforms import (
+    TransformChange,
+    TransformResult,
+    TransformStatus,
+)
+from codeshift.migrator.llm_migrator import LLMMigrator
+
+
+class MigrationEngine:
+    """Orchestrates migrations using a tiered approach.
+
+    Tier 1: Deterministic AST transforms for well-known libraries
+    Tier 2: Knowledge base guided migration with LLM assistance
+    Tier 3: Pure LLM migration for unknown patterns
+    """
+
+    def __init__(
+        self,
+        llm_migrator: LLMMigrator | None = None,
+    ):
+        """Initialize the migration engine.
+
+        Args:
+            llm_migrator: Optional LLM migrator instance.
+        """
+        self.llm_migrator = llm_migrator or LLMMigrator()
+
+    def run_migration(
+        self,
+        code: str,
+        file_path: Path,
+        library: str,
+        old_version: str,
+        new_version: str,
+        knowledge_base: GeneratedKnowledgeBase | None = None,
+        progress_callback: Callable[[str], None] | None = None,
+    ) -> TransformResult:
+        """Run migration using the appropriate tier.
+
+        Args:
+            code: Source code to migrate.
+            file_path: Path to the file being migrated.
+            library: Library being upgraded.
+            old_version: Current version.
+            new_version: Target version.
+            knowledge_base: Optional generated knowledge base.
+            progress_callback: Optional progress callback.
+
+        Returns:
+            TransformResult with migrated code.
+        """
+
+        def report(msg: str) -> None:
+            if progress_callback:
+                progress_callback(msg)
+
+        # Tier 1: Try deterministic AST transforms for known libraries
+        if is_tier_1_library(library):
+            report(f"Using Tier 1 (deterministic AST transforms) for {library}")
+            result = self._apply_tier1_transform(code, file_path, library)
+
+            if result.status == TransformStatus.SUCCESS:
+                return result
+
+            # If partial success, try Tier 2/3 for remaining
+            if result.status == TransformStatus.PARTIAL:
+                report("Tier 1 partial - attempting Tier 2/3 for remaining changes")
+                # Use transformed code as base for next tier
+                code = result.transformed_code
+
+        # Tier 2: Knowledge base guided migration
+        if knowledge_base and knowledge_base.overall_confidence >= Confidence.MEDIUM:
+            report("Using Tier 2 (KB-guided migration)")
+            result = self._apply_tier2_transform(
+                code,
+                file_path,
+                library,
+                old_version,
+                new_version,
+                knowledge_base,
+            )
+
+            if result.status == TransformStatus.SUCCESS:
+                return result
+
+        # Tier 3: Pure LLM migration
+        if self.llm_migrator.is_available:
+            report("Using Tier 3 (LLM-assisted migration)")
+            return self._apply_tier3_transform(
+                code,
+                file_path,
+                library,
+                old_version,
+                new_version,
+            )
+
+        # No migration possible
+        return TransformResult(
+            file_path=file_path,
+            status=TransformStatus.NO_CHANGES,
+            original_code=code,
+            transformed_code=code,
+            errors=["No migration method available"],
+        )
+
+    def _apply_tier1_transform(
+        self,
+        code: str,
+        file_path: Path,
+        library: str,
+    ) -> TransformResult:
+        """Apply Tier 1 deterministic AST transforms.
+
+        Args:
+            code: Source code to transform.
+            file_path: Path to the file.
+            library: Library being upgraded.
+
+        Returns:
+            TransformResult.
+        """
+        # Import transformers dynamically based on library
+        transform_func = self._get_transform_func(library)
+
+        if transform_func is None:
+            return TransformResult(
+                file_path=file_path,
+                status=TransformStatus.NO_CHANGES,
+                original_code=code,
+                transformed_code=code,
+                errors=[f"No Tier 1 transformer for {library}"],
+            )
+
+        try:
+            transformed_code, changes = transform_func(code)
+
+            return TransformResult(
+                file_path=file_path,
+                status=TransformStatus.SUCCESS if changes else TransformStatus.NO_CHANGES,
+                original_code=code,
+                transformed_code=transformed_code,
+                changes=[
+                    TransformChange(
+                        description=c.description,
+                        line_number=c.line_number,
+                        original=c.original,
+                        replacement=c.replacement,
+                        transform_name=c.transform_name,
+                        confidence=getattr(c, "confidence", 1.0),
+                    )
+                    for c in changes
+                ],
+            )
+        except Exception as e:
+            return TransformResult(
+                file_path=file_path,
+                status=TransformStatus.FAILED,
+                original_code=code,
+                transformed_code=code,
+                errors=[f"Tier 1 transform failed: {e}"],
+            )
+
+    def _apply_tier2_transform(
+        self,
+        code: str,
+        file_path: Path,
+        library: str,
+        old_version: str,
+        new_version: str,
+        knowledge_base: GeneratedKnowledgeBase,
+    ) -> TransformResult:
+        """Apply Tier 2 knowledge base guided migration.
+
+        Args:
+            code: Source code to transform.
+            file_path: Path to the file.
+            library: Library being upgraded.
+            old_version: Current version.
+            new_version: Target version.
+            knowledge_base: Generated knowledge base.
+
+        Returns:
+            TransformResult.
+        """
+        if not self.llm_migrator.is_available:
+            return TransformResult(
+                file_path=file_path,
+                status=TransformStatus.NO_CHANGES,
+                original_code=code,
+                transformed_code=code,
+                errors=["LLM not available for Tier 2"],
+            )
+
+        # Build context from knowledge base
+        context_parts = [
+            f"Breaking changes for {library} {old_version} -> {new_version}:",
+        ]
+
+        for change in knowledge_base.breaking_changes:
+            if change.new_api:
+                context_parts.append(
+                    f"- {change.old_api} -> {change.new_api}: {change.description}"
+                )
+            else:
+                context_parts.append(f"- {change.old_api} (removed): {change.description}")
+
+        context = "\n".join(context_parts)
+
+        # Use LLM with context
+        result = self.llm_migrator.migrate(
+            code=code,
+            library=library,
+            from_version=old_version,
+            to_version=new_version,
+            context=context,
+        )
+
+        if result.success:
+            return TransformResult(
+                file_path=file_path,
+                status=TransformStatus.SUCCESS,
+                original_code=code,
+                transformed_code=result.migrated_code,
+                changes=[
+                    TransformChange(
+                        description="KB-guided LLM migration",
+                        line_number=1,
+                        original="(various)",
+                        replacement="(migrated)",
+                        transform_name="tier2_kb_guided",
+                        confidence=0.9,
+                    )
+                ],
+            )
+
+        return TransformResult(
+            file_path=file_path,
+            status=TransformStatus.FAILED,
+            original_code=code,
+            transformed_code=code,
+            errors=[result.error or "Tier 2 migration failed"],
+        )
+
+    def _apply_tier3_transform(
+        self,
+        code: str,
+        file_path: Path,
+        library: str,
+        old_version: str,
+        new_version: str,
+    ) -> TransformResult:
+        """Apply Tier 3 pure LLM migration.
+
+        Args:
+            code: Source code to transform.
+            file_path: Path to the file.
+            library: Library being upgraded.
+            old_version: Current version.
+            new_version: Target version.
+
+        Returns:
+            TransformResult.
+        """
+        result = self.llm_migrator.migrate(
+            code=code,
+            library=library,
+            from_version=old_version,
+            to_version=new_version,
+        )
+
+        if result.success:
+            return TransformResult(
+                file_path=file_path,
+                status=TransformStatus.SUCCESS,
+                original_code=code,
+                transformed_code=result.migrated_code,
+                changes=[
+                    TransformChange(
+                        description="LLM-assisted migration",
+                        line_number=1,
+                        original="(various)",
+                        replacement="(migrated)",
+                        transform_name="tier3_llm",
+                        confidence=0.7,
+                        notes="Review carefully - LLM-generated changes",
+                    )
+                ],
+            )
+
+        return TransformResult(
+            file_path=file_path,
+            status=TransformStatus.NO_CHANGES,
+            original_code=code,
+            transformed_code=code,
+            errors=[result.error or "Tier 3 migration failed"],
+        )
+
+    def _get_transform_func(self, library: str) -> Callable | None:
+        """Get the transform function for a library.
+
+        Args:
+            library: Library name.
+
+        Returns:
+            Transform function or None.
+        """
+        try:
+            if library == "pydantic":
+                from codeshift.migrator.transforms.pydantic_v1_to_v2 import (
+                    transform_pydantic_v1_to_v2,
+                )
+
+                return transform_pydantic_v1_to_v2
+            elif library == "fastapi":
+                from codeshift.migrator.transforms.fastapi_transformer import (
+                    transform_fastapi,
+                )
+
+                return transform_fastapi
+            elif library == "sqlalchemy":
+                from codeshift.migrator.transforms.sqlalchemy_transformer import (
+                    transform_sqlalchemy,
+                )
+
+                return transform_sqlalchemy
+            elif library == "pandas":
+                from codeshift.migrator.transforms.pandas_transformer import (
+                    transform_pandas,
+                )
+
+                return transform_pandas
+            elif library == "requests":
+                from codeshift.migrator.transforms.requests_transformer import (
+                    transform_requests,
+                )
+
+                return transform_requests
+        except ImportError:
+            pass
+
+        return None
+
+
+# Singleton instance
+_default_engine: MigrationEngine | None = None
+
+
+def get_migration_engine() -> MigrationEngine:
+    """Get the default migration engine instance."""
+    global _default_engine
+    if _default_engine is None:
+        _default_engine = MigrationEngine()
+    return _default_engine
+
+
+def run_migration(
+    code: str,
+    file_path: Path,
+    library: str,
+    old_version: str,
+    new_version: str,
+    knowledge_base: GeneratedKnowledgeBase | None = None,
+) -> TransformResult:
+    """Convenience function to run a migration.
+
+    Args:
+        code: Source code to migrate.
+        file_path: Path to the file.
+        library: Library being upgraded.
+        old_version: Current version.
+        new_version: Target version.
+        knowledge_base: Optional generated knowledge base.
+
+    Returns:
+        TransformResult.
+    """
+    engine = get_migration_engine()
+    return engine.run_migration(
+        code=code,
+        file_path=file_path,
+        library=library,
+        old_version=old_version,
+        new_version=new_version,
+        knowledge_base=knowledge_base,
+    )