Statements and Code Blocks

Codegen uses two classes to represent code structure at the highest level:

• Statement: Represents a single line or block of code

  • Can be assignments, imports, loops, conditionals, etc.
  • Contains source code, dependencies, and type information
  • May contain nested code blocks (like in functions or loops)

• CodeBlock: A container for multiple Statements

  • Found in files, functions, classes, and control flow blocks
  • Provides APIs for analyzing and manipulating statements
  • Handles scope, variables, and dependencies

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.

See Adding Statements for details.

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:

Import Statements / Export Statements

See imports and exports for more details.

# 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

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.

Switch/Match Statements

# 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 Statements

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}")

Assignment Statements

# 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]}")