feat: RunCodemodTool (#557)

Author: jayhack

src/codegen/extensions/langchain/tools.py

@@ -20,6 +20,7 @@
 from codegen.extensions.tools.link_annotation import add_links_to_message
 from codegen.extensions.tools.replacement_edit import replacement_edit
 from codegen.extensions.tools.reveal_symbol import reveal_symbol
+from codegen.extensions.tools.run_codemod import run_codemod
 from codegen.extensions.tools.search import search
 from codegen.extensions.tools.semantic_edit import semantic_edit
 from codegen.extensions.tools.semantic_search import semantic_search
@@ -709,7 +710,8 @@ def get_workspace_tools(codebase: Codebase) -> list["BaseTool"]:
         RenameFileTool(codebase),
         ReplacementEditTool(codebase),
         RevealSymbolTool(codebase),
-        RunBashCommandTool(),  # Note: This tool doesn't need the codebase
+        RunBashCommandTool(),
+        RunCodemodTool(codebase),
         SearchTool(codebase),
         SemanticEditTool(codebase),
         SemanticSearchTool(codebase),
@@ -770,3 +772,38 @@ def _run(
             count=count,
         )
         return json.dumps(result, indent=2)
+
+
+class RunCodemodInput(BaseModel):
+    """Input for running a codemod."""
+
+    codemod_source: str = Field(
+        ...,
+        description="""Source code of the codemod function. Must define a 'run(codebase: Codebase)' function that makes the desired changes.
+Example:
+```python
+def run(codebase: Codebase):
+    for file in codebase.files:
+        if file.filepath.endswith('.py'):
+            content = file.content
+            # Make changes to content
+            file.edit(new_content)
+```
+""",
+    )
+
+
+class RunCodemodTool(BaseTool):
+    """Tool for running custom codemod functions."""
+
+    name: ClassVar[str] = "run_codemod"
+    description: ClassVar[str] = "Run a custom codemod function to make systematic changes across the codebase"
+    args_schema: ClassVar[type[BaseModel]] = RunCodemodInput
+    codebase: Codebase = Field(exclude=True)
+
+    def __init__(self, codebase: Codebase) -> None:
+        super().__init__(codebase=codebase)
+
+    def _run(self, codemod_source: str) -> str:
+        result = run_codemod(self.codebase, codemod_source)
+        return json.dumps(result, indent=2)

src/codegen/extensions/tools/__init__.py

@@ -19,6 +19,7 @@
 from .rename_file import rename_file
 from .replacement_edit import replacement_edit
 from .reveal_symbol import reveal_symbol
+from .run_codemod import run_codemod
 from .search import search
 from .semantic_edit import semantic_edit
 from .semantic_search import semantic_search
@@ -45,6 +46,7 @@
     "rename_file",
     "replacement_edit",
     "reveal_symbol",
+    "run_codemod",
     # Search operations
     "search",
     # Edit operations

src/codegen/extensions/tools/run_codemod.py

@@ -0,0 +1,71 @@
+"""Tool for running custom codemod functions on the codebase."""
+
+import importlib.util
+import sys
+from pathlib import Path
+from tempfile import NamedTemporaryFile
+from typing import Any
+
+from codegen import Codebase
+
+
+def run_codemod(codebase: Codebase, codemod_source: str) -> dict[str, Any]:
+    """Run a custom codemod function on the codebase.
+
+    The codemod_source should define a function like:
+    ```python
+    def run(codebase: Codebase):
+        # Make changes to the codebase
+        ...
+    ```
+
+    Args:
+        codebase: The codebase to operate on
+        codemod_source: Source code of the codemod function
+
+    Returns:
+        Dict containing execution results and diffs
+
+    Raises:
+        ValueError: If codemod source is invalid or execution fails
+    """
+    # Create a temporary module to run the codemod
+    with NamedTemporaryFile(suffix=".py", mode="w", delete=False) as temp_file:
+        # Add imports and write the codemod source
+        temp_file.write("from codegen import Codebase\n\n")
+        temp_file.write(codemod_source)
+        temp_file.flush()
+
+        try:
+            # Import the temporary module
+            spec = importlib.util.spec_from_file_location("codemod", temp_file.name)
+            if not spec or not spec.loader:
+                msg = "Failed to create module spec"
+                raise ValueError(msg)
+
+            module = importlib.util.module_from_spec(spec)
+            sys.modules["codemod"] = module
+            spec.loader.exec_module(module)
+
+            # Verify run function exists
+            if not hasattr(module, "run"):
+                msg = "Codemod must define a 'run' function"
+                raise ValueError(msg)
+
+            # Run the codemod
+            module.run(codebase)
+            codebase.commit()
+            diff = codebase.get_diff()
+
+            return {
+                "status": "success",
+                "diff": diff,
+            }
+
+        except Exception as e:
+            msg = f"Codemod execution failed: {e!s}"
+            raise ValueError(msg)
+
+        finally:
+            # Clean up temporary file
+            Path(temp_file.name).unlink()

tests/unit/codegen/extensions/test_tools.py

@@ -14,6 +14,7 @@
     rename_file,
     replacement_edit,
     reveal_symbol,
+    run_codemod,
     search,
     semantic_edit,
     semantic_search,
@@ -240,3 +241,16 @@ def test_replacement_edit(codebase):
     )
     assert result["status"] == "unchanged"
     assert "No matches found" in result["message"]
+
+
+def test_run_codemod(codebase):
+    """Test running custom codemods."""
+    # Test adding type hints
+    codemod_source = """
+def run(codebase: Codebase):
+    for file in codebase.files:
+        file.edit('# hello, world!' + file.content)
+"""
+    result = run_codemod(codebase, codemod_source)
+    assert result["status"] == "success"
+    assert "+# hello, world" in result["diff"]