Increase Type Coverage

Learn how to manipulate type annotations in Codegen’s GraphSitter library.

This guide demonstrates how to analyze and manipulate type annotations with GraphSitter.

Common use cases include:

• Adding a type to a union or generic type
• Checking if a generic type has a given subtype
• Resolving a type annotation

Finding the extent of your type coverage

First, to get an indication of your progress on type coverage, count the number of function arguments that have type hints in your codebase: //: # (—CountTypedParametersSkill—) //: # (—THE CODE BELOW IS AUTO GENERATED. UPDATE THE SNIPPET BY UPDATING THE SKILL—)

You can perform an analogous operation for return types, method parameters, and more. Note that the code presented here only analyzes non-method functions.

Adding simple return type annotations

To add a return type, use function.set_return_type. The script below will add a -> None return type to all functions that contain no return statements: //: # (—SetReturnTypeToNoneForFunctionsWithoutReturns—) //: # (—THE CODE BELOW IS AUTO GENERATED. UPDATE THE SNIPPET BY UPDATING THE SKILL—)

Adding a type to a union type

Since a union type is a collection, you can call standard methods such as append to add new types to it.

Adding a type to a generic type

You can add parameters to a generic type in a similar fashion //: # (—ModifyReturnTypeWithNewParameter—) //: # (—THE CODE BELOW IS AUTO GENERATED. UPDATE THE SNIPPET BY UPDATING THE SKILL—)

function = codebase.get_file("path/to/file.py").get_function("function_name")
# def function_name() -> tuple[a, b]: ...

# Add a new parameter to the return type
function.return_type.parameters.append("float")

Checking if a generic type contains a given type

You can also check if a given type is in the list of parameters

function = codebase.get_file("path/to/file.py").get_function("function_name")
# def function_name() -> tuple[a, b]: ...

# Check if "a" is in the function's return_type's parameters
if "a" in function.return_type.parameters:
    # type "a" is present in the return type parameters
    print("type 'a' is present in the return type parameters.")

Resolving a type_annotation

You can find the referenced symbol of any type //: # (—InspectResolveFunctionReturnType—) //: # (—THE CODE BELOW IS AUTO GENERATED. UPDATE THE SNIPPET BY UPDATING THE SKILL—)

Resolving a type_annotation of an assignment

You can also use resolved_type directly on assignments to get the resolved_symbol

a = codebase.get_file("path/to/file.py").get_global_var("a")
# a: MyType = ...
print(a.type)  # Resolves to MyType
a.type.rename("MyNewType")  # Renames the symbol `MyType` to `MyNewType` throughout the codebase