Codegen provides comprehensive APIs for working with function calls through several key classes:
- FunctionCall - Represents a function invocation
- Argument - Represents arguments passed to a function
- Parameter - Represents parameters in a function definition
Navigating Function Calls
Codegen provides two main ways to navigate function calls:
- From a function to its call sites using call_sites
- From a function to the calls it makes (within it’s CodeBlock) using function_calls
Here’s how to analyze function usage patterns:
# Find the most called function
most_called = max(codebase.functions, key=lambda f: len(f.call_sites))
print(f"\nMost called function: {most_called.name}")
print(f"Called {len(most_called.call_sites)} times from:")
for call in most_called.call_sites:
print(f" - {call.parent_function.name} at line {call.start_point[0]}")
# Find function that makes the most calls
most_calls = max(codebase.functions, key=lambda f: len(f.function_calls))
print(f"\nFunction making most calls: {most_calls.name}")
print(f"Makes {len(most_calls.function_calls)} calls to:")
for call in most_calls.function_calls:
print(f" - {call.name}")
# Find functions with no callers (potential dead code)
unused = [f for f in codebase.functions if len(f.call_sites) == 0]
print(f"\nUnused functions:")
for func in unused:
print(f" - {func.name} in {func.filepath}")
# Find recursive functions
recursive = [f for f in codebase.functions
if any(call.name == f.name for call in f.function_calls)]
print(f"\nRecursive functions:")
for func in recursive:
print(f" - {func.name}")
This navigation allows you to:
- Find heavily used functions
- Analyze call patterns
- Map dependencies between functions
Arguments and Parameters
The Argument class represents values passed to a function, while Parameter represents the receiving variables in the function definition:
Consider the following code:
# Source code:
def process_data(input_data: str, debug: bool = False):
pass
process_data("test", debug=True)
You can access and modify the arguments and parameters of the function call with APIs detailed below.
Finding Arguments
The primary APIs for finding arguments are:
# Get the function call
call = file.function_calls[0]
# Working with arguments
for arg in call.args:
print(f"Arg {arg.index}: {arg.value}") # Access argument value
print(f"Is named: {arg.is_named}") # Check if it's a kwarg
print(f"Name: {arg.name}") # For kwargs, e.g. "debug"
# Get corresponding parameter
if param := arg.parameter:
print(f"Parameter type: {param.type}")
print(f"Is optional: {param.is_optional}")
print(f"Has default: {param.default}")
# Finding specific arguments
debug_arg = call.get_arg_by_parameter_name("debug")
first_arg = call.get_arg_by_index(0)
Modifying Arguments
There are two ways to modify function call arguments:
- Using FunctionCall.set_kwarg(…) to add or modify keyword arguments:
# Modifying keyword arguments
call.set_kwarg("debug", "False") # Modifies existing kwarg
call.set_kwarg("new_param", "value", create_on_missing=True) # Adds new kwarg
call.set_kwarg("input_data", "'new_value'", override_existing=True) # Converts positional to kwarg
- Using FuncionCall.args.append(…) to add new arguments:
# Adding new arguments
call.args.append('cloud="aws"') # Add a new keyword argument
call.args.append('"value"') # Add a new positional argument
# Real-world example: Adding arguments to a decorator
@app.function(image=runner_image)
def my_func():
pass
# Add cloud and region if not present
if "cloud=" not in decorator.call.source:
decorator.call.args.append('cloud="aws"')
if "region=" not in decorator.call.source:
decorator.call.args.append('region="us-east-1"')
The set_kwarg
method provides intelligent argument manipulation:
- If the argument exists and is positional, it converts it to a keyword argument
- If the argument exists and is already a keyword, it updates its value (if override_existing=True)
- If the argument doesn’t exist, it creates it (if create_on_missing=True)
- When creating new arguments, it intelligently places them based on parameter order
Arguments and parameters support safe edit operations like so:
# Modifying arguments
debug_arg.edit("False") # Change argument value
first_arg.add_keyword("input_data") # Convert to named argument
# modifying parameters
param = codebase.get_function('process_data').get_parameter('debug')
param.rename('_debug') # updates all call-sites
param.set_type_annotation('bool')
Finding Function Definitions
Every FunctionCall can navigate to its definition through function_definition and function_definitions:
function_call = codebase.files[0].function_calls[0]
function_definition = function_call.function_definition
print(f"Definition found in: {function_definition.filepath}")
Finding Parent (Containing) Functions
FunctionCalls can access the function that invokes it via parent_function.
For example, given the following code:
# Source code:
def outer():
def inner():
helper()
inner()
You can find the parent function of the helper call:
# Manipulation code:
# Find the helper() call
helper_call = file.get_function("outer").function_calls[1]
# Get containing function
parent = helper_call.parent_function
print(f"Call is inside: {parent.name}") # 'inner'
# Get the full call hierarchy
outer = parent.parent_function
print(f"Which is inside: {outer.name}") # 'outer'
Method Chaining
Codegen enables working with chained method calls through predecessor and related properties:
For example, for the following database query:
# Source code:
query.select(Table)
.where(id=1)
.order_by("name")
.limit(10)
You can access the chain of calls:
# Manipulation code:
# Get the `limit` call in the chain
limit_call = next(f for f in file.function.function_calls if f.name == "limit", None)
# Navigate backwards through the chain
order_by = limit_call.predecessor
where = order_by.predecessor
select = where.predecessor
# Get the full chain at once
chain = limit_call.call_chain # [select, where, order_by, limit]
# Access the root object
base = limit_call.base # Returns the 'query' object
# Check call relationships
print(f"After {order_by.name}: {limit_call.name}")
print(f"Before {where.name}: {select.name}")