Codegen uses two classes to represent code structure at the highest level:
Codegen provides rich APIs for working with code statements and blocks, allowing you to analyze and manipulate code structure at a granular level.
Working with Statements
Basic Usage
Every file, function, and class in Codegen has a CodeBlock that contains its statements:
# Access statements in a file
file = codebase.get_file("main.py")
for statement in file.code_block.statements:
print(f"Statement type: {statement.statement_type}")
# Access statements in a function
function = file.get_function("process_data")
for statement in function.code_block.statements:
print(f"Statement: {statement.source}")
Filtering Statements
Filter through statements using Python’s builtin isinstance function.
# Filter statements by type
for stmt in file.code_block.statements:
if isinstance(stmt, ImportStatement):
print(stmt)
Adding Statements
Functions and Files support .prepend_statement(…) and .add_statement(…) to add statements to the symbol.
Working with Nested Structures
Frequently you will want to check if a statement is nested within another structure, for example if a statement is inside an if block or a try/catch statement.
Codegen supports this functionality with the Editable.is_wrapped_in(…) method.
func = codebase.get_function("process_data")
for usage in func.local_variable_usages:
if usage.is_wrapped_in(IfStatement):
print(f"Usage of {usage.name} is inside an if block")
Similarly, all Editable objects support the .parent_statement, which can be used to navigate the statement hierarchy.
func = codebase.get_function("process_data")
for usage in func.local_variable_usages:
if isinstance(usage.parent_statement, IfStatement):
print(f"Usage of {usage.name} is directly beneath an IfStatement")
Wrapping and Unwrapping Statements
CodeBlocks support wrapping and unwrapping with the following APIs:
- .wrap(…) - allows you to wrap a statement in a new structure.
- .unwrap(…) - allows you to remove the wrapping structure while preserving the code block’s contents.
# Wrap code blocks with new structures
function.code_block.wrap("with open('test.txt', 'w') as f:")
# Result:
# with open('test.txt', 'w') as f:
# original_code_here...
# Wrap code in a function
file.code_block.wrap("def process_data(a, b):")
# Result:
# def process_data(a, b):
# original_code_here...
# Unwrap code from its container
if_block.code_block.unwrap() # Removes the if statement but keeps its body
while_loop.code_block.unwrap() # Removes the while loop but keeps its body
Both wrap and unwrap are potentially unsafe changes and will modify
business logic.
The unwrap() method preserves the indentation of the code block’s contents
while removing the wrapping structure. This is useful for refactoring nested
code structures.
Statement Types
Codegen supports various statement types, each with specific APIs:
# Access import statements
for import_stmt in file.import_statements:
print(f"Module: {import_stmt.module}")
for imported in import_stmt.imports:
print(f" Imported: {imported.name}")
# Remove specific imports
import_stmt = file.import_statements[0]
import_stmt.imports[0].remove() # Remove first import
# Remove entire import statement
import_stmt.remove()
If/Else statements provide rich APIs for analyzing and manipulating conditional logic:
# Access if/else blocks
if_block = file.code_block.statements[0]
print(f"Condition: {if_block.condition.source}")
# Check block types
if if_block.is_if_statement:
print("Main if block")
elif if_block.is_elif_statement:
print("Elif block")
elif if_block.is_else_statement:
print("Else block")
# Access alternative blocks
for elif_block in if_block.elif_statements:
print(f"Elif condition: {elif_block.condition.source}")
if else_block := if_block.else_statement:
print("Has else block")
# Access nested code blocks
for block in if_block.nested_code_blocks:
print(f"Block statements: {len(block.statements)}")
If blocks also support condition reduction, which can simplify conditional logic:
# Reduce if condition to True
if_block.reduce_condition(True)
# Before:
# if condition:
# print("a")
# else:
# print("b")
# After:
# print("a")
# Reduce elif condition to False
elif_block.reduce_condition(False)
# Before:
# if a:
# print("a")
# elif condition:
# print("b")
# else:
# print("c")
# After:
# if a:
# print("a")
# else:
# print("c")
When reducing conditions, Codegen automatically handles the restructuring of
elif/else chains and preserves the correct control flow.
# TypeScript switch statements
switch_stmt = file.code_block.statements[0]
for case_stmt in switch_stmt.cases:
print(f"Case condition: {case_stmt.condition}")
print(f"Is default: {case_stmt.default}")
# Access statements in each case
for statement in case_stmt.code_block.statements:
print(f"Statement: {statement.source}")
# Python match statements
match_stmt = file.code_block.statements[0]
for case in match_stmt.cases:
print(f"Pattern: {case.pattern}")
for statement in case.code_block.statements:
print(f"Statement: {statement.source}")
while_stmt = file.code_block.statements[0]
print(f"Condition: {while_stmt.condition}")
# Access loop body
for statement in while_stmt.code_block.statements:
print(f"Body statement: {statement.source}")
# Get function calls within the loop
for call in while_stmt.function_calls:
print(f"Function call: {call.source}")
# Access assignments in a code block
for statement in code_block.statements:
if statement.statement_type == StatementType.ASSIGNMENT:
for assignment in statement.assignments:
print(f"Variable: {assignment.name}")
print(f"Value: {assignment.value}")
Working with Code Blocks
Code blocks provide several ways to analyze and manipulate their content:
Statement Access
code_block = function.code_block
# Get all statements
all_statements = code_block.statements
# Get statements by type
if_blocks = code_block.if_blocks
while_loops = code_block.while_loops
try_blocks = code_block.try_blocks
# Get local variables
local_vars = code_block.get_local_var_assignments()
Statement Dependencies
# Get dependencies between statements
function = file.get_function("process")
for statement in function.code_block.statements:
deps = statement.dependencies
print(f"Statement {statement.source} depends on: {[d.name for d in deps]}")
Parent-Child Relationships
# Access parent statements
function = file.get_function("main")
parent_stmt = function.parent_statement
# Access nested symbols
class_def = file.get_class("MyClass")
for method in class_def.methods:
parent = method.parent_statement
print(f"Method {method.name} is defined in {parent.source}")
Common Operations
Finding Statements
# Find specific statements
assignments = [s for s in code_block.statements
if s.statement_type == StatementType.ASSIGNMENT]
# Find statements by content
matching = [s for s in code_block.statements
if "specific_function()" in s.source]
Analyzing Flow Control
# Analyze control flow
for statement in code_block.statements:
if statement.statement_type == StatementType.IF_BLOCK:
print("Condition:", statement.condition)
print("Then:", statement.consequence_block.statements)
if statement.alternative_block:
print("Else:", statement.alternative_block.statements)
Working with Functions
# Analyze function calls in statements
for statement in code_block.statements:
for call in statement.function_calls:
print(f"Calls function: {call.name}")
print(f"With arguments: {[arg.source for arg in call.arguments]}")
