Reusable Codemods

Codegen enables you to create reusable code transformations using Python functions decorated with @codegen.function. These codemods can be shared, versioned, and run by your team.

Creating Codemods

The easiest way to create a new codemod is using the CLI create command:

codegen create rename-function .

This creates a new codemod in your .codegen/codemods directory:

import codegen
from codegen import Codebase

@codegen.function("rename-function")
def run(codebase: Codebase):
    """Add a description of what this codemod does."""
    # Add your code here
    pass

Codemods are stored in .codegen/codemods/name/name.py and are tracked in Git for easy sharing.

AI-Powered Generation with -d

You can use AI to generate an initial implementation by providing a description:

codegen create rename-function . -d "Rename the getUserData function to fetchUserProfile"

This will:

1. Generate an implementation based on your description
2. Create a custom system prompt that you can provide to an IDE chat assistant (learn more about working with AI)
3. Place both files in the codemod directory

Running Codemods

Once created, run your codemod using:

codegen run rename-function

The execution flow:

1. Codegen parses your codebase into a graph representation
2. Your codemod function is executed against this graph
3. Changes are tracked and applied to your filesystem
4. A diff preview shows what changed

Codemod Structure

A codemod consists of three main parts:

1. The @codegen.function decorator that names your codemod
2. A run function that takes a Codebase parameter
3. Your transformation logic using the Codebase API

import codegen
from codegen import Codebase

@codegen.function("update-imports")
def run(codebase: Codebase):
    """Update import statements to use new package names."""
    for file in codebase.files:
        for imp in file.imports:
            if imp.module == "old_package":
                imp.rename("new_package")
    codebase.commit()

Arguments

Codemods can accept arguments using Pydantic models:

from pydantic import BaseModel

class RenameArgs(BaseModel):
    old_name: str
    new_name: str

@codegen.function("rename-function")
def run(codebase: Codebase, arguments: RenameArgs):
    """Rename a function across the codebase."""
    old_func = codebase.get_function(arguments.old_name)
    if old_func:
        old_func.rename(arguments.new_name)
    codebase.commit()

Run it with:

codegen run rename-function --arguments '{"old_name": "getUserData", "new_name": "fetchUserProfile"}'

Directory Structure

Your codemods live in a dedicated directory structure:

.codegen/
└── codemods/
    └── rename_function/
        ├── rename_function.py       # The codemod implementation
        └── rename_function_prompt.md  # System prompt (if using AI)