run now works

Author: codegen-bot

src/codegen/cli/commands/run/main.py

@@ -1,6 +1,5 @@
 import json
 
-import rich
 import rich_click as click
 
 from codegen.cli.auth.session import CodegenSession
@@ -13,59 +12,35 @@
 @requires_init
 @click.argument("label", required=True)
 @click.option("--web", is_flag=True, help="Run the function on the web service instead of locally")
-@click.option("--apply-local", is_flag=True, help="Applies the generated diff to the repository")
 @click.option("--diff-preview", type=int, help="Show a preview of the first N lines of the diff")
 @click.option("--arguments", type=str, help="Arguments as a json string to pass as the function's 'arguments' parameter")
 def run_command(
     session: CodegenSession,
     label: str,
     web: bool = False,
-    apply_local: bool = False,
     diff_preview: int | None = None,
     arguments: str | None = None,
 ):
     """Run a codegen function by its label."""
-    # First try to find it as a stored codemod
-    codemod = CodemodManager.get(label)
-    if codemod:
-        if codemod.arguments_type_schema and not arguments:
-            raise click.ClickException(f"This function requires the --arguments parameter. Expected schema: {codemod.arguments_type_schema}")
+    # Get and validate the codemod
+    codemod = CodemodManager.get_codemod(label)
 
-        if codemod.arguments_type_schema and arguments:
-            arguments_json = json.loads(arguments)
-            is_valid = validate_json(codemod.arguments_type_schema, arguments_json)
-            if not is_valid:
-                raise click.ClickException(f"Invalid arguments format. Expected schema: {codemod.arguments_type_schema}")
+    # Handle arguments if needed
+    if codemod.arguments_type_schema and not arguments:
+        raise click.ClickException(f"This function requires the --arguments parameter. Expected schema: {codemod.arguments_type_schema}")
 
-        if web:
-            from codegen.cli.commands.run.run_cloud import run_cloud
-
-            run_cloud(session, codemod, apply_local=apply_local, diff_preview=diff_preview)
-        else:
-            from codegen.cli.commands.run.run_local import run_local
-
-            run_local(session, codemod, apply_local=apply_local, diff_preview=diff_preview)
-        return
-
-    # If not found as a stored codemod, look for decorated functions
-    functions = CodemodManager.get_decorated()
-    matching = [f for f in functions if f.name == label]
-
-    if not matching:
-        raise click.ClickException(f"No function found with label '{label}'")
-
-    if len(matching) > 1:
-        # If multiple matches, show their locations
-        rich.print(f"[yellow]Multiple functions found with label '{label}':[/yellow]")
-        for func in matching:
-            rich.print(f"  • {func.filepath}")
-        raise click.ClickException("Please specify the exact file with codegen run <path>")
+    if codemod.arguments_type_schema and arguments:
+        arguments_json = json.loads(arguments)
+        is_valid = validate_json(codemod.arguments_type_schema, arguments_json)
+        if not is_valid:
+            raise click.ClickException(f"Invalid arguments format. Expected schema: {codemod.arguments_type_schema}")
 
+    # Run the codemod
     if web:
         from codegen.cli.commands.run.run_cloud import run_cloud
 
-        run_cloud(session, matching[0], apply_local=apply_local, diff_preview=diff_preview)
+        run_cloud(session, codemod, diff_preview=diff_preview)
     else:
         from codegen.cli.commands.run.run_local import run_local
 
-        run_local(session, matching[0], apply_local=apply_local, diff_preview=diff_preview)
+        run_local(session, codemod, diff_preview=diff_preview)

src/codegen/cli/commands/run/run_local.py

@@ -2,88 +2,70 @@
 
 import rich
 from rich.panel import Panel
+from rich.status import Status
 
 from codegen import Codebase
 from codegen.cli.auth.session import CodegenSession
-from codegen.cli.git.patch import apply_patch
-from codegen.cli.rich.codeblocks import format_command
-from codegen.cli.rich.spinners import create_spinner
 from codegen.cli.utils.function_finder import DecoratedFunction
 
 
+def parse_codebase(repo_root: Path) -> Codebase:
+    """Parse the codebase at the given root.
+
+    Args:
+        repo_root: Path to the repository root
+
+    Returns:
+        Parsed Codebase object
+    """
+    codebase = Codebase(repo_root)
+    return codebase
+
+
 def run_local(
     session: CodegenSession,
     function: DecoratedFunction,
-    apply_local: bool = False,
     diff_preview: int | None = None,
 ) -> None:
     """Run a function locally against the codebase.
 
     Args:
         session: The current codegen session
-        function: The function to run (either a DecoratedFunction or Codemod)
-        apply_local: Whether to apply changes to the local filesystem
+        function: The function to run
         diff_preview: Number of lines of diff to preview (None for all)
     """
-    # Initialize codebase from git repo root
+    # Parse codebase and run
     repo_root = Path(session.git_repo.workdir)
 
-    with create_spinner("Parsing codebase...") as status:
-        codebase = Codebase(repo_root)
-
-    try:
-        # Run the function
-        rich.print(f"Running {function.name} locally...")
-        result = function.run(codebase)
-
-        if not result:
-            rich.print("\n[yellow]No changes were produced by this codemod[/yellow]")
-            return
-
-        # Show diff preview if requested
-        if diff_preview:
-            rich.print("")  # Add spacing
-            diff_lines = result.splitlines()
-            truncated = len(diff_lines) > diff_preview
-            limited_diff = "\n".join(diff_lines[:diff_preview])
-
-            if truncated:
-                if apply_local:
-                    limited_diff += f"\n\n...\n\n[yellow]diff truncated to {diff_preview} lines, view the full change set in your local file system[/yellow]"
-                else:
-                    limited_diff += f"\n\n...\n\n[yellow]diff truncated to {diff_preview} lines, view the full change set by running with --apply-local[/yellow]"
-
-            panel = Panel(limited_diff, title="[bold]Diff Preview[/bold]", border_style="blue", padding=(1, 2), expand=False)
-            rich.print(panel)
-
-        # Apply changes if requested
-        if apply_local:
-            try:
-                apply_patch(session.git_repo, f"\n{result}\n")
-                rich.print("")
-                rich.print("[green]✓ Changes have been applied to your local filesystem[/green]")
-                rich.print("[yellow]→ Don't forget to commit your changes:[/yellow]")
-                rich.print(format_command("git add ."))
-                rich.print(format_command("git commit -m 'Applied codemod changes'"))
-            except Exception as e:
-                rich.print("")
-                rich.print("[red]✗ Failed to apply changes locally[/red]")
-                rich.print("\n[yellow]This usually happens when you have uncommitted changes.[/yellow]")
-                rich.print("\nOption 1 - Save your changes:")
-                rich.print("  1. [blue]git status[/blue]        (check your working directory)")
-                rich.print("  2. [blue]git add .[/blue]         (stage your changes)")
-                rich.print("  3. [blue]git commit -m 'msg'[/blue]  (commit your changes)")
-                rich.print("  4. Run this command again")
-                rich.print("\nOption 2 - Discard your changes:")
-                rich.print("  1. [red]git reset --hard HEAD[/red]     (⚠️ discards all uncommitted changes)")
-                rich.print("  2. [red]git clean -fd[/red]            (⚠️ removes all untracked files)")
-                rich.print("  3. Run this command again\n")
-                raise RuntimeError("Failed to apply patch to local filesystem") from e
-        else:
-            rich.print("")
-            rich.print("To apply these changes locally:")
-            rich.print(format_command(f"codegen run {function.name} --apply-local"))
-
-    except Exception as e:
-        rich.print(f"[red]Error running {function.name}:[/red] {e!s}")
-        raise
+    with Status("[bold]Parsing codebase...", spinner="dots") as status:
+        codebase = parse_codebase(repo_root)
+        status.update("[bold green]✓ Parsed codebase")
+
+        status.update("[bold]Running codemod...")
+        function.run(codebase)  # Run the function
+        status.update("[bold green]✓ Completed codemod")
+
+    # Get the diff from the codebase
+    result = codebase.get_diff()
+
+    # Handle no changes case
+    if not result:
+        rich.print("\n[yellow]No changes were produced by this codemod[/yellow]")
+        return
+
+    # Show diff preview if requested
+    if diff_preview:
+        rich.print("")  # Add spacing
+        diff_lines = result.splitlines()
+        truncated = len(diff_lines) > diff_preview
+        limited_diff = "\n".join(diff_lines[:diff_preview])
+
+        if truncated:
+            limited_diff += f"\n\n...\n\n[yellow]diff truncated to {diff_preview} lines[/yellow]"
+
+        panel = Panel(limited_diff, title="[bold]Diff Preview[/bold]", border_style="blue", padding=(1, 2), expand=False)
+        rich.print(panel)
+
+    # Apply changes
+    rich.print("")
+    rich.print("[green]✓ Changes have been applied to your local filesystem[/green]")

src/codegen/cli/sdk/decorator.py

@@ -36,6 +36,8 @@ def __call__(self, func: Callable[P, T]) -> Callable[P, T]:
         def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
             return func(*args, **kwargs)
 
+        # Set the codegen name on the wrapper function
+        wrapper.__codegen_name__ = self.name
         self.func = wrapper
         return wrapper
 

src/codegen/cli/utils/codemod_manager.py

@@ -1,6 +1,8 @@
 import builtins
 from pathlib import Path
 
+import rich_click as click
+
 from codegen.cli.utils.function_finder import DecoratedFunction, find_codegen_functions
 
 
@@ -25,6 +27,39 @@ class CodemodManager:
     def get_valid_name(name: str) -> str:
         return name.lower().replace(" ", "_").replace("-", "_")
 
+    @classmethod
+    def get_codemod(cls, name: str, start_path: Path | None = None) -> DecoratedFunction:
+        """Get and validate a codemod by name.
+
+        Args:
+            name: Name of the codemod to find
+            start_path: Directory to start searching from (default: current directory)
+
+        Returns:
+            The validated DecoratedFunction
+
+        Raises:
+            click.ClickException: If codemod can't be found or loaded
+        """
+        # First try to find the codemod
+        codemod = cls.get(name, start_path)
+        if not codemod:
+            # If not found, check if any codemods exist
+            all_codemods = cls.list(start_path)
+            if not all_codemods:
+                raise click.ClickException("No codemods found. Create one with:\n" + "  codegen create my-codemod")
+            else:
+                available = "\n  ".join(f"- {c.name}" for c in all_codemods)
+                raise click.ClickException(f"Codemod '{name}' not found. Available codemods:\n  {available}")
+
+        # Verify we can import it
+        try:
+            # This will raise ValueError if function can't be imported
+            codemod.validate()
+            return codemod
+        except Exception as e:
+            raise click.ClickException(f"Error loading codemod '{name}': {e!s}")
+
     @classmethod
     def list(cls, start_path: Path | None = None) -> builtins.list[DecoratedFunction]:
         """List all codegen decorated functions in Python files under the given path.

src/codegen/cli/utils/function_finder.py

@@ -18,6 +18,60 @@ class DecoratedFunction:
     parameters: list[tuple[str, str | None]] = dataclasses.field(default_factory=list)
     arguments_type_schema: dict | None = None
 
+    def run(self, codebase) -> str | None:
+        """Import and run the actual function from its file.
+
+        Args:
+            codebase: The codebase to run the function on
+
+        Returns:
+            The result of running the function (usually a diff string)
+        """
+        if not self.filepath:
+            raise ValueError("Cannot run function without filepath")
+
+        # Import the module containing the function
+        spec = importlib.util.spec_from_file_location("module", self.filepath)
+        if not spec or not spec.loader:
+            raise ImportError(f"Could not load module from {self.filepath}")
+
+        module = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(module)
+
+        # Find the decorated function
+        for item_name in dir(module):
+            item = getattr(module, item_name)
+            if hasattr(item, "__codegen_name__") and item.__codegen_name__ == self.name:
+                # Found our function, run it
+                return item(codebase)
+
+        raise ValueError(f"Could not find function '{self.name}' in {self.filepath}")
+
+    def validate(self) -> None:
+        """Verify that this function can be imported and accessed.
+
+        Raises:
+            ValueError: If the function can't be found or imported
+        """
+        if not self.filepath:
+            raise ValueError("Cannot validate function without filepath")
+
+        # Import the module containing the function
+        spec = importlib.util.spec_from_file_location("module", self.filepath)
+        if not spec or not spec.loader:
+            raise ImportError(f"Could not load module from {self.filepath}")
+
+        module = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(module)
+
+        # Find the decorated function
+        for item_name in dir(module):
+            item = getattr(module, item_name)
+            if hasattr(item, "__codegen_name__") and item.__codegen_name__ == self.name:
+                return  # Found it!
+
+        raise ValueError(f"Could not find function '{self.name}' in {self.filepath}")
+
 
 class CodegenFunctionVisitor(ast.NodeVisitor):
     def __init__(self):