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:
- Relationship tracking
- 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
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")