Manipulate Function Call Sites

Learn how to manipulate function call sites with Codegen

Codegen’s GraphSitter library provides powerful APIs for systematically manipulating function call sites throughout your codebase.

This guide will walk you through the basics of working with function call sites using GraphSitter.

Common use cases include:

• Updating function calls after changing function signatures
• Refactoring to use named arguments
• Adding or removing arguments from function calls

Overview

To manipulate function call sites, you first need to locate them. You can do this using the call_sites property of a function.

# Get the specific function we want to modify
function = codebase.get_file("path/to/file.py").get_function("my_function")

# Iterate through all places where this function is called
for call_site in function.call_sites:
    # Add the new keyword argument
    call_site.set_kwarg("new_arg", "42")

    # Process each existing argument
    for arg in call_site.args:
        if not arg.is_named:
            # Convert positional args to named args for manipulation
            arg.add_keyword(arg.name)

        # Remove specified argument
        if arg.name == "existing_arg":
            arg.remove()
        # Rename argument
        elif arg.name == "arg_to_rename":
            arg.add_keyword("new_arg_name")
        # Modify argument value
        elif arg.name == "arg_to_modify":
            arg.edit("42")

Modifying Arguments

You can add, remove, or modify arguments at a call site:

# Add new argument with value 42
call_site.set_kwarg("new_arg", "42")

# Remove old_arg if it exists
if arg := call_site.get_arg_by_parameter_name("old_arg"):
    arg.remove()

# Modify existing_arg if it exists
if arg := call_site.get_arg_by_parameter_name("existing_arg"):
    arg.edit("existing_arg='new_value'")

Converting to Named Arguments

You can convert positional arguments to named arguments:

call_site.convert_args_to_kwargs()

This will convert all positional arguments to named arguments, based on the parameter names in the function definition.

Updating Multiple Call Sites

If you need to update multiple call sites, you can do so in a loop:

function = codebase.get_file("path/to/file.py").get_function("my_function")

# Update all call sites of this function
for call in function.call_sites:
    call.set_kwarg("new_arg", "42")

Handling Different Argument Types

GraphSitter allows you to handle different types of arguments:

for arg in call_site.args:
    if arg.is_named:
        print(f"  Found named arg in {file.filepath}: {arg.name} = {arg.value}")
    else:
        print(f"  Found positional arg in {file.filepath}: {arg.value}")

Best Practices

1. Commit After Major Changes: If you’re making multiple significant changes, use codebase.commit() between them to ensure the codebase graph is up-to-date.

2. Re-fetch References: After a commit, re-fetch any file or symbol references you’re working with, as they may have become stale.

3. Handle Errors: Be prepared to handle cases where functions or call sites might not exist, or where modifications might fail due to syntax errors.

By following these guidelines, you can effectively manipulate function call sites throughout your codebase while maintaining its integrity and functionality.