feat: codegen init creates + perists .codegen/.venv (#124)

jayhack · codegen-bot · codegen-bot · commit 23d702bef5fa · 2025-01-28T03:56:28.000+01:00
# Motivation
&lt;!-- Why is this change necessary? --&gt;

# Content
&lt;!-- Please include a summary of the change --&gt;
# Testing
&lt;!-- How was the change tested? --&gt;
# Please check the following before marking your PR as ready for review

- [ ] I have added tests for my changes
- [ ] I have updated the documentation or added new documentation as
needed
- [ ] I have read and agree to the [Contributor License
Agreement](../CLA.md)

---------

Co-authored-by: codegen-bot &lt;team+codegenbot@codegen.sh&gt;
diff --git a/docs/building-with-codegen/dot-codegen.mdx b/docs/building-with-codegen/dot-codegen.mdx
@@ -11,12 +11,11 @@ The `.codegen` directory contains your project's Codegen configuration, codemods
 
 ```bash
 .codegen/
-├── config.toml         # Project configuration
-├── codemods/          # Your codemod implementations
-├── jupyter/           # Jupyter notebooks for exploration
-├── docs/             # API documentation
-├── examples/         # Example code
-└── prompts/          # AI system prompts
+├── .venv/            # Python virtual environment (gitignored)
+├── config.toml       # Project configuration
+├── codemods/         # Your codemod implementations
+├── jupyter/          # Jupyter notebooks for exploration
+└── codegen-system-prompt.txt  # AI system prompt
 ```
 
 ## Initialization
@@ -31,6 +30,20 @@ codegen init [--fetch-docs] [--repo-name NAME] [--organization-name ORG]
 The `--fetch-docs` flag downloads API documentation and examples specific to your project's programming language.
 </Note>
 
+## Virtual Environment
+
+Codegen maintains its own virtual environment in `.codegen/.venv/` to ensure consistent package versions and isolation from your project's dependencies. This environment is:
+
+- Created using `uv` for fast, reliable package management
+- Initialized with Python 3.13
+- Automatically managed by Codegen commands
+- Used for running codemods and Jupyter notebooks
+- Gitignored to avoid committing environment-specific files
+
+The environment is created during `codegen init` and used by commands like `codegen run` and `codegen notebook`.
+
+<Note>To debug codemods, you will need to set the python virtual environment in your IDE to `.codegen/.venv`</Note>
+
 ### Configuration
 
 The `config.toml` file stores your project settings:
@@ -49,13 +62,15 @@ Codegen automatically adds appropriate entries to your `.gitignore`:
 
 ```gitignore
 # Codegen
-.codegen/prompts/
+.codegen/.venv/
 .codegen/docs/
-.codegen/examples/
+.codegen/jupyter/
+.codegen/codegen-system-prompt.txt
 ```
 
 <Info>
-While prompts, docs, and examples are ignored, your codemods in `.codegen/codemods/` are tracked in Git.
+- While most directories are ignored, your codemods in `.codegen/codemods/` and `config.toml` are tracked in Git
+- The virtual environment and Jupyter notebooks are gitignored to avoid environment-specific issues
 </Info>
 
 ## Working with Codemods
diff --git a/docs/introduction/overview.mdx b/docs/introduction/overview.mdx
@@ -9,7 +9,7 @@ iconType: "solid"
 
 It provides a scriptable interface to a powerful, multi-lingual language server built on top of [Tree-sitter](https://tree-sitter.github.io/tree-sitter/).
 
-export const intoSnippet = `# Codegen builds a complete graph connecting
+export const metaCode = `# Codegen builds a complete graph connecting
 #  functions, classes, imports and their relationships
 from codegen import Codebase
 
@@ -21,13 +21,23 @@ for function in codebase.functions:
         function.remove()
 `
 
+export const code = `def foo():
+  pass
+
+def bar():
+  foo()
+
+def baz():
+  pass
+`
+
 <iframe
   width="100%"
   height="300px"
   scrolling="no"
-  src={`https://codegen.sh/embedded/codemod/?code=${encodeURIComponent(
-    intoSnippet
-  )}`}
+  src={`https://chadcode.sh/embedded/codemod/?code=${encodeURIComponent(
+    metaCode
+  )}&input=${encodeURIComponent(code)}`}
   style={{
     backgroundColor: "#15141b",
   }}
diff --git a/src/codegen/cli/commands/create/main.py b/src/codegen/cli/commands/create/main.py
@@ -12,6 +12,7 @@
 from codegen.cli.rich.pretty_print import pretty_print_error
 from codegen.cli.rich.spinners import create_spinner
 from codegen.cli.utils.constants import ProgrammingLanguage
+from codegen.cli.utils.default_code import DEFAULT_CODEMOD
 from codegen.cli.workspace.decorators import requires_init
 
 
@@ -62,19 +63,6 @@ def make_relative(path: Path) -> str:
         return f"./{path.name}"
 
 
-def get_default_code(name: str) -> str:
-    """Get the default function code without using the API."""
-    return f'''import codegen
-from codegen import Codebase
-
-@codegen.function("{name}")
-def run(codebase: Codebase):
-    """Add a description of what this codemod does."""
-    # Add your code here
-    pass
-'''
-
-
 @click.command(name="create")
 @requires_init
 @click.argument("name", type=str)
@@ -111,7 +99,7 @@ def create_command(session: CodegenSession, name: str, path: Path, description:
                     prompt_path.write_text(response.context)
         else:
             # Use default implementation
-            code = get_default_code(name)
+            code = DEFAULT_CODEMOD.format(name=name)
 
         # Create the target directory if needed
         target_path.parent.mkdir(parents=True, exist_ok=True)
diff --git a/src/codegen/cli/commands/init/render.py b/src/codegen/cli/commands/init/render.py
@@ -4,6 +4,7 @@
 def get_success_message(codegen_dir: Path, docs_dir: Path, examples_dir: Path) -> str:
     """Get the success message to display after initialization."""
     return """📁 .codegen configuration folder created:
-   [dim]config.toml[/dim]  Project configuration
-   [dim]codemods/[/dim]    Your codemod implementations
-   [dim]codegen-system-prompt.txt[/dim]     AI system prompt (gitignored)"""
+   [dim]config.toml[/dim]                Project configuration
+   [dim]codemods/[/dim]                  Your codemod implementations
+   [dim].venv/[/dim]                     Python virtual environment (gitignored)
+   [dim]codegen-system-prompt.txt[/dim]  AI system prompt (gitignored)"""
diff --git a/src/codegen/cli/commands/notebook/main.py b/src/codegen/cli/commands/notebook/main.py
@@ -8,6 +8,7 @@
 from codegen.cli.auth.decorators import requires_auth
 from codegen.cli.auth.session import CodegenSession
 from codegen.cli.rich.spinners import create_spinner
+from codegen.cli.utils.notebooks import create_notebook
 from codegen.cli.workspace.decorators import requires_init
 from codegen.cli.workspace.venv_manager import VenvManager
 
@@ -19,57 +20,24 @@ def create_jupyter_dir() -> Path:
     return jupyter_dir
 
 
-def create_notebook(jupyter_dir: Path) -> Path:
-    """Create a new Jupyter notebook if it doesn't exist."""
-    notebook_path = jupyter_dir / "tmp.ipynb"
-    if not notebook_path.exists():
-        notebook_content = {
-            "cells": [
-                {
-                    "cell_type": "code",
-                    "execution_count": None,
-                    "metadata": {},
-                    "outputs": [],
-                    "source": ["from codegen import Codebase\n", "\n", "# Initialize codebase\n", "codebase = Codebase('../../')\n"],
-                }
-            ],
-            "metadata": {"kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"}},
-            "nbformat": 4,
-            "nbformat_minor": 4,
-        }
-        import json
-
-        notebook_path.write_text(json.dumps(notebook_content, indent=2))
-    return notebook_path
-
-
 @click.command(name="notebook")
-@click.option("--background", is_flag=True, help="Run Jupyter Lab in the background")
 @requires_auth
 @requires_init
-def notebook_command(session: CodegenSession, background: bool = False):
+def notebook_command(session: CodegenSession):
     """Open a Jupyter notebook with the current codebase loaded."""
     with create_spinner("Setting up Jupyter environment...") as status:
         venv = VenvManager()
 
-        if not venv.is_initialized():
-            status.update("Creating virtual environment...")
-            venv.create_venv()
-
-        status.update("Installing required packages...")
-        venv.install_packages("codegen", "jupyterlab")
+        status.update("Checking Jupyter installation...")
+        venv.ensure_jupyter()
 
         jupyter_dir = create_jupyter_dir()
         notebook_path = create_notebook(jupyter_dir)
 
-        status.update("Starting Jupyter Lab...")
+        status.update("Running Jupyter Lab...")
 
         # Prepare the environment with the virtual environment activated
         env = {**os.environ, "VIRTUAL_ENV": str(venv.venv_dir), "PATH": f"{venv.venv_dir}/bin:{os.environ['PATH']}"}
 
-        # Start Jupyter Lab
-        if background:
-            subprocess.Popen(["jupyter", "lab", str(notebook_path)], env=env, start_new_session=True)
-        else:
-            # Run in foreground
-            subprocess.run(["jupyter", "lab", str(notebook_path)], env=env, check=True)
+        # Run Jupyter Lab
+        subprocess.run(["jupyter", "lab", str(notebook_path)], env=env, check=True)
diff --git a/src/codegen/cli/commands/run/main.py b/src/codegen/cli/commands/run/main.py
@@ -1,11 +1,13 @@
 import json
+import os
 
 import rich_click as click
 
 from codegen.cli.auth.session import CodegenSession
 from codegen.cli.utils.codemod_manager import CodemodManager
 from codegen.cli.utils.json_schema import validate_json
 from codegen.cli.workspace.decorators import requires_init
+from codegen.cli.workspace.venv_manager import VenvManager
 
 
 @click.command(name="run")
@@ -22,6 +24,15 @@ def run_command(
     arguments: str | None = None,
 ):
     """Run a codegen function by its label."""
+    # Ensure venv is initialized
+    venv = VenvManager()
+    if not venv.is_initialized():
+        raise click.ClickException("Virtual environment not found. Please run 'codegen init' first.")
+
+    # Set up environment with venv
+    os.environ["VIRTUAL_ENV"] = str(venv.venv_dir)
+    os.environ["PATH"] = f"{venv.venv_dir}/bin:{os.environ['PATH']}"
+
     # Get and validate the codemod
     codemod = CodemodManager.get_codemod(label)
 
diff --git a/src/codegen/cli/utils/default_code.py b/src/codegen/cli/utils/default_code.py
@@ -0,0 +1,20 @@
+DEFAULT_CODEMOD = '''import codegen
+from codegen import Codebase
+
+
+@codegen.function("{name}")
+def run(codebase: Codebase):
+    """Add a description of what this codemod does."""
+    # Add your code here
+    print('Total files: ', len(codebase.files))
+    print('Total functions: ', len(codebase.functions))
+    print('Total imports: ', len(codebase.imports))
+
+
+if __name__ == "__main__":
+    print('Parsing codebase...')
+    codebase = Codebase("./")
+
+    print('Running...')
+    run(codebase)
+'''
diff --git a/src/codegen/cli/utils/notebooks.py b/src/codegen/cli/utils/notebooks.py
@@ -0,0 +1,44 @@
+import json
+from pathlib import Path
+
+DEFAULT_NOTEBOOK_CONTENT = """from codegen import Codebase
+
+# Initialize codebase
+codebase = Codebase('../../')
+
+# Print out stats
+print("🔍 Codebase Analysis")
+print("=" * 50)
+print(f"📚 Total Files: {len(codebase.files)}")
+print(f"⚡ Total Functions: {len(codebase.functions)}")
+print(f"🔄 Total Imports: {len(codebase.imports)}")
+""".strip()
+
+
+def create_notebook(jupyter_dir: Path) -> Path:
+    """Create a new Jupyter notebook if it doesn't exist.
+
+    Args:
+        jupyter_dir: Directory where the notebook should be created
+
+    Returns:
+        Path to the created or existing notebook
+    """
+    notebook_path = jupyter_dir / "tmp.ipynb"
+    if not notebook_path.exists():
+        notebook_content = {
+            "cells": [
+                {
+                    "cell_type": "code",
+                    "execution_count": None,
+                    "metadata": {},
+                    "outputs": [],
+                    "source": DEFAULT_NOTEBOOK_CONTENT,
+                }
+            ],
+            "metadata": {"kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"}},
+            "nbformat": 4,
+            "nbformat_minor": 4,
+        }
+        notebook_path.write_text(json.dumps(notebook_content, indent=2))
+    return notebook_path
diff --git a/src/codegen/cli/workspace/initialize_workspace.py b/src/codegen/cli/workspace/initialize_workspace.py
@@ -13,46 +13,10 @@
 from codegen.cli.git.repo import get_git_repo
 from codegen.cli.git.url import get_git_organization_and_repo
 from codegen.cli.rich.spinners import create_spinner
+from codegen.cli.utils.notebooks import create_notebook
 from codegen.cli.workspace.docs_workspace import populate_api_docs
 from codegen.cli.workspace.examples_workspace import populate_examples
-
-DEFAULT_CODE = """
-from codegen import Codebase
-
-# Initialize codebase
-codebase = Codebase('../../')
-
-# Print out stats
-print("🔍 Codebase Analysis")
-print("=" * 50)
-print(f"📚 Total Files: {len(codebase.files)}")
-print(f"⚡ Total Functions: {len(codebase.functions)}")
-print(f"🔄 Total Imports: {len(codebase.imports)}")
-"""
-
-
-def create_notebook(jupyter_dir: Path) -> Path:
-    """Create a new Jupyter notebook if it doesn't exist."""
-    notebook_path = jupyter_dir / "tmp.ipynb"
-    if not notebook_path.exists():
-        notebook_content = {
-            "cells": [
-                {
-                    "cell_type": "code",
-                    "execution_count": None,
-                    "metadata": {},
-                    "outputs": [],
-                    "source": [DEFAULT_CODE],
-                }
-            ],
-            "metadata": {"kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"}},
-            "nbformat": 4,
-            "nbformat_minor": 4,
-        }
-        import json
-
-        notebook_path.write_text(json.dumps(notebook_content, indent=2))
-    return notebook_path
+from codegen.cli.workspace.venv_manager import VenvManager
 
 
 def initialize_codegen(
@@ -93,6 +57,13 @@ def initialize_codegen(
         JUPYTER_DIR.mkdir(parents=True, exist_ok=True)
         CODEMODS_DIR.mkdir(parents=True, exist_ok=True)
 
+        # Initialize virtual environment
+        status_obj.update(f"   {'Creating' if isinstance(status, str) else 'Checking'} virtual environment...")
+        venv = VenvManager()
+        if not venv.is_initialized():
+            venv.create_venv()
+            venv.install_packages("codegen")
+
         # Download system prompt
         try:
             from codegen.cli.api.endpoints import CODEGEN_SYSTEM_PROMPT_URL
@@ -164,7 +135,8 @@ def modify_gitignore(codegen_folder: Path):
         "examples/",
         "prompts/",
         "jupyter/",
-        "codegen-system-prompt.txt",  # Add system prompt to gitignore
+        ".venv/",  # Add venv to gitignore
+        "codegen-system-prompt.txt",
         "",
         "# Python cache files",
         "__pycache__/",
diff --git a/src/codegen/cli/workspace/venv_manager.py b/src/codegen/cli/workspace/venv_manager.py
