Add memory profiling and optimization to CLI

src/codegen/cli/cli.py

@@ -11,6 +11,7 @@
 from codegen.cli.commands.login.main import login_command
 from codegen.cli.commands.logout.main import logout_command
 from codegen.cli.commands.lsp.lsp import lsp_command
+from codegen.cli.commands.memprof.main import memprof_command
 from codegen.cli.commands.notebook.main import notebook_command
 from codegen.cli.commands.profile.main import profile_command
 from codegen.cli.commands.reset.main import reset_command
@@ -51,6 +52,7 @@
 main.add_command(lsp_command)
 main.add_command(serve_command)
 main.add_command(start_command)
+main.add_command(memprof_command)  # Add the memory profiling command
 
 
 if __name__ == "__main__":

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

@@ -0,0 +1,69 @@
+import os
+from typing import Optional
+
+import rich
+import rich_click as click
+from rich import box
+from rich.panel import Panel
+
+from codegen.cli.utils.memory_profiler import profile_command
+
+
+@click.command(name="memprof")
+@click.argument("command", nargs=-1, required=True)
+@click.option(
+    "--output-dir",
+    "-o",
+    type=click.Path(file_okay=False),
+    help="Directory to save memory profile reports",
+)
+def memprof_command(command: list[str], output_dir: Optional[str] = None):
+    """Profile memory usage of a Codegen CLI command.
+
+    Example:
+        codegen memprof run my-codemod --arguments '{"param": "value"}'
+    """
+    if not command:
+        rich.print("[bold red]Error:[/bold red] No command specified")
+        return
+
+    # Convert command tuple to list
+    cmd_args = list(command)
+
+    # Set default output directory if not provided
+    if not output_dir:
+        home_dir = os.path.expanduser("~")
+        output_dir = os.path.join(home_dir, ".codegen", "memory_profiles")
+
+    # Run the profiling
+    rich.print(
+        Panel(
+            f"[cyan]Profiling command:[/cyan] codegen {' '.join(cmd_args)}",
+            title="🔍 [bold]Memory Profiler[/bold]",
+            border_style="cyan",
+            box=box.ROUNDED,
+            padding=(1, 2),
+        )
+    )
+
+    try:
+        report_dir = profile_command(cmd_args, output_dir=output_dir)
+        rich.print(
+            Panel(
+                f"[green]Memory profile saved to:[/green] {report_dir}",
+                title="✅ [bold]Profiling Complete[/bold]",
+                border_style="green",
+                box=box.ROUNDED,
+                padding=(1, 2),
+            )
+        )
+    except Exception as e:
+        rich.print(
+            Panel(
+                f"[red]Error during profiling:[/red] {e!s}",
+                title="❌ [bold]Profiling Failed[/bold]",
+                border_style="red",
+                box=box.ROUNDED,
+                padding=(1, 2),
+            )
+        )

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

@@ -1,5 +1,9 @@
+import gc
+import os
+import time
 from pathlib import Path
 
+import psutil
 import rich
 from rich.panel import Panel
 from rich.status import Status
@@ -27,12 +31,15 @@
     Returns:
         Parsed Codebase object
     """
+    # Force garbage collection before parsing to free up memory
+    gc.collect()
+
     codebase = Codebase(
         projects=[
             ProjectConfig(
                 repo_operator=RepoOperator(repo_config=RepoConfig.from_repo_path(repo_path=repo_path)),
                 subdirectories=subdirectories,
                 programming_language=language or determine_project_language(repo_path),
             )
         ]
     )
@@ -51,21 +58,36 @@
         function: The function to run
         diff_preview: Number of lines of diff to preview (None for all)
     """
+    # Get initial memory usage
+    process = psutil.Process(os.getpid())
+    initial_memory = process.memory_info().rss / (1024 * 1024)  # Convert to MB
+
     # Parse codebase and run
     with Status(f"[bold]Parsing codebase at {session.repo_path} with subdirectories {function.subdirectories or 'ALL'} and language {function.language or 'AUTO'} ...", spinner="dots") as status:
+        start_time = time.time()
         codebase = parse_codebase(repo_path=session.repo_path, subdirectories=function.subdirectories, language=function.language)
-        status.update("[bold green]✓ Parsed codebase")
+        parse_time = time.time() - start_time
+        status.update(f"[bold green]✓ Parsed codebase in {parse_time:.2f}s")
+
+        # Memory usage after parsing
+        post_parse_memory = process.memory_info().rss / (1024 * 1024)
 
         status.update("[bold]Running codemod...")
+        start_time = time.time()
         function.run(codebase)  # Run the function
-        status.update("[bold green]✓ Completed codemod")
+        run_time = time.time() - start_time
+        status.update(f"[bold green]✓ Completed codemod in {run_time:.2f}s")
 
     # Get the diff from the codebase
     result = codebase.get_diff()
 
+    # Final memory usage
+    final_memory = process.memory_info().rss / (1024 * 1024)
+
     # Handle no changes case
     if not result:
         rich.print("\n[yellow]No changes were produced by this codemod[/yellow]")
+        rich.print(f"\n[dim]Memory usage: {initial_memory:.2f}MB → {final_memory:.2f}MB (Δ {final_memory - initial_memory:.2f}MB)[/dim]")
         return
 
     # Show diff preview if requested
@@ -84,3 +106,11 @@
     # Apply changes
     rich.print("")
     rich.print("[green]✓ Changes have been applied to your local filesystem[/green]")
+
+    # Print memory usage statistics
+    rich.print(f"\n[dim]Memory usage: {initial_memory:.2f}MB → {final_memory:.2f}MB (Δ {final_memory - initial_memory:.2f}MB)[/dim]")
+    rich.print(f"[dim]Parsing: {parse_time:.2f}s, Execution: {run_time:.2f}s[/dim]")
+
+    # Clean up to free memory
+    del codebase
+    gc.collect()