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

See Migrating APIs for relevant tutorials and applications.

Codegen provides two main ways to navigate function calls:

  1. From a function to its call sites using call_sites
  2. 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:

  1. 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
  1. Using FuncionCall.args.append(…) to add new arguments:

    FunctionCall.args is a Collection of Argument objects, so it supports .append(…), .insert(…) and other collection methods.

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

Was this page helpful?

Dependencies and UsagesVariable Assignments