Manipulating datastructures

Codegen’s GraphSitter library parses all collections into structure similar to the built-in python datastructures they represent. This guide will walk you through the process modifying and using these structures.

Common use cases include:

• Adding parameters to functions

• Converting dictionaries to schemas

• Appending to lists of arguments

​

Adding parameters to functions

Since parameters are a collection, you can call standard methods such as append to add new parameters to a function.

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

# Add a new parameter to the function
function.parameters.append("new_param: int")

​

Inspecting the contents of a Dictionary

You can get the value of a dictionary to use in further modifications.

var = codebase.get_file("path/to/file.py").get_global_var("foo")

# Get the value of the dictionary for key "example_key"
example = var.value["example_key"]
assert example == "example_value"

# Alternatively
example = var.value.get("example_key", None)
assert example == "example_value"

# Set the value of the dictionary for key "example_key" to "xyz"
var.value["example_key"] = '"xyz"'

​

Converting dictionaries to schemas

You can convert dictionaries to schemas by modifying the value of a variable.

var = codebase.get_file("path/to/file.py").get_global_var("var_name")

# Convert its assignment to a Schema and excludes certain keys
var.set_value(f"Schema({', '.join(f'{k}={v}' for k, v in var.value.items() if k != 'excluded')})")

​

Appending to a list of in an assignment

You can append to a list using standard methods. For example, to append to a list assigned to a global variable

var = codebase.get_file("path/to/file.py").get_global_var("foo")

# Assert the type is List (a GraphSitter type)
if not isinstance(var.value, List):
    raise ValueError(f"Expected a list, but found {type(var.value)}")

# Append to the list
var.value.append('"bar"')

​

Check if a collection contains an item

You can use python’s in operator to check if a collection contains a specific item. It checks if any contained item has the same source.

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

# Check if the function has a decorator named "@cache"
if "@cache" in function.decorators:
    print(f"Function: {function.name} has a @cache decorator")