The Editable API

Every code element in Codegen is an Editable - meaning it can be manipulated while maintaining correctness.

All higher-level code manipulation APIs are built on top of the atomic Editable API.

Core Concepts

Every Editable provides:

• Information about the source code:
  • source - the text content of the Editable
  • extended_source - includes relevant content like decorators, comments, etc.
• Information about the file that contains the Editable:
  • file - the SourceFile that contains this Editable
• Relationship tracking
  • parent_class - the Class that contains this Editable
  • parent_function - the Function that contains this Editable
  • parent_statement - the Statement that contains this Editable
• Safe modification operations

Basic Editing

There are several fundamental ways to modify code with Editables:

# 1. edit() - Replace entire source with new content
function = codebase.get_function("process_data")
function.edit("""
def process_data(input_data: dict) -> dict:
    return transform(input_data)
""")

# 2. Replace - Substitute text while preserving context
class_def = codebase.get_class("UserModel")
class_def.replace("user_id", "account_id")  # Updates all occurrences

# 3. Remove - Safely delete code with proper cleanup
unused_import = file.get_import("from utils import deprecated_func")
unused_import.remove()  # Handles formatting, commas, etc

# 4. Insert - Add code before or after an element
function.insert_before("# Process user input")  # Adds comment before function
function.insert_after("""
def validate_data(data: dict) -> bool:
    return all(required in data for required in REQUIRED_FIELDS)
""")  # Adds new function after

Finding and Searching

Editables provide powerful search capabilities:

# Find string literals
results = function.find_string_literals(["error", "warning"])
results = function.find_string_literals(["error"], fuzzy_match=True)

# Search with regex
matches = function.search(r"data\['[^']*'\]")  # Find dict access
matches = function.search("TODO:", include_comments=True)

# Find specific patterns
variables = function.get_variable_usages("config")
function_calls = function.function_calls  # All function calls within this node

Smart Formatting

Codegen handles formatting details automatically:

# Adding to import statements
import_stmt = file.get_import("from mylib import func1")
import_stmt.add_symbol("func2")  # Handles comma placement
import_stmt.add_symbol("func3")  # Maintains proper formatting

# Multi-line formatting is preserved
from mylib import (
    func1,
    func2,    # New imports maintain
    func3     # existing style
)

Safe Removals

Removing code elements is safe and clean:

# Remove a function and its decorators
function.remove()  # Removes associated comments and formatting

# Remove imports cleanly
import_stmt.remove()  # Handles commas and whitespace

Working with References

Editables track their relationships to other code elements:

# Find and update all references
function = codebase.get_function("old_name")
function.rename("new_name")  # Updates all usages

# Navigate relationships
print(function.parent_function)  # Containing function
print(function.parent_class)    # Containing class
print(function.parent_statement)  # Containing statement

Understanding Context

Editables provide rich information about their location and context in the code:

Parent Relationships

# Get containing elements
function = codebase.get_function("process_data")
print(function.parent_class)     # Class containing this function
print(function.parent_function)  # Function containing this function (for nested functions)
print(function.parent_statement) # Statement containing this function

# Check if top-level
is_top_level = function.parent_class is None and function.parent_function is None

Statement Containment

The is_wrapped_in method lets you check if an Editable is contained within specific types of statements:

# Check containment in statement types
is_in_try = function.is_wrapped_in("try")
is_in_if = function.is_wrapped_in("if")
is_in_while = function.is_wrapped_in("while")

# Get the first parent statements of a certain type
if_block = function.parent_of_type(IfStatement)

# Common patterns
if function.is_wrapped_in(IfStatement):
    print("This is in an IfBlock")

if variable.is_wrapped_in(WithStatement):
    print("Variable used in WithStatement")

Common Use Cases

# Move nested functions to module level
for func in file.functions:
    if func.parent_function:  # This is a nested function
        func.parent_function.insert_before(func.source) # Move to module level
        func.remove() # Remove the nested function

# Find variables defined in unsafe blocks
for var in function.code_block.get_local_var_assignments():
    if var.is_wrapped_in(TryStatement):
        print(f"Warning: {var.name} defined in try block")