This guide covers the core APIs and patterns for working with type annotations in Codegen.

​
Type Resolution

Codegen builds a complete dependency graph of your codebase, connecting functions, classes, imports, and their relationships. This enables powerful type resolution capabilities:

from codegen import Codebase

# Initialize codebase with dependency graph
codebase = Codebase("./")

# Get a function with a type annotation
function = codebase.get_file("path/to/file.py").get_function("my_func")

# Resolve its return type to actual symbols
return_type = function.return_type
resolved_symbols = return_type.resolved_types  # Returns the actual Symbol objects

# For generic types, you can resolve parameters
if hasattr(return_type, "parameters"):
    for param in return_type.parameters:
        resolved_param = param.resolved_types  # Get the actual type parameter symbols

# For assignments, resolve their type
assignment = codebase.get_file("path/to/file.py").get_assignment("my_var")
resolved_type = assignment.type.resolved_types

Type resolution follows imports and handles complex cases like type aliases, forward references, and generic type parameters.

​
Core Interfaces

Type annotations in Codegen are built on two key interfaces:

  • Typeable - The base interface for any node that can have a type annotation (parameters, variables, functions, etc). Provides .type and .is_typed.
  • Type - The base class for all type annotations. Provides type resolution and dependency tracking.

Any node that inherits from Typeable will have a .type property that returns a Type object, which can be used to inspect and modify type annotations.

Learn more about inheritable behaviors like Typeable here

​
Core Type APIs

Type annotations can be accessed and modified through several key APIs:

​
Function Types

The main APIs for function types are Function.return_type and Function.set_return_type:

# Get return type
return_type = function.return_type  # -> TypeAnnotation
print(return_type.source)  # "List[str]"
print(return_type.is_typed)  # True/False

# Set return type
function.set_return_type("List[str]")
function.set_return_type(None)  # Removes type annotation

​
Parameter Types

Parameters use Parameter.type and Parameter.set_type_annotation:

for param in function.parameters:
    # Get parameter type
    param_type = param.type  # -> TypeAnnotation
    print(param_type.source)  # "int"
    print(param_type.is_typed)  # True/False

    # Set parameter type
    param.set_type("int")
    param.set_type(None)  # Removes type annotation

​
Variable Types

Variables and attributes use Assignment.type and Assignment.set_type_annotation. This applies to:

# For global/local assignments
assignment = file.get_assignment("my_var")
var_type = assignment.type  # -> TypeAnnotation
print(var_type.source)  # "str"

# Set variable type
assignment.set_type("str")
assignment.set_type(None)  # Removes type annotation

# For class attributes
class_def = file.get_class("MyClass")
for attr in class_def.attributes:
    # Each attribute has an assignment property
    attr_type = attr.assignment.type  # -> TypeAnnotation
    print(f"{attr.name}: {attr_type.source}")  # e.g. "x: int"
    
    # Set attribute type
    attr.assignment.set_type("int")

# You can also access attributes directly by index
first_attr = class_def.attributes[0]
first_attr.assignment.set_type("str")

​
Working with Complex Types

​
Union Types

Union types (UnionType) can be manipulated as collections:

# Get union type
union_type = function.return_type  # -> A | B 
print(union_type.symbols)  # ["A", "B"]

# Add/remove options
union_type.append("float")
union_type.remove("None")

# Check contents
if "str" in union_type.options:
    print("String is a possible type")
Learn more about working with collections here

​
Generic Types

Generic types (GenericType) expose their parameters as collection of Parameters:

# Get generic type
generic_type = function.return_type  # -> GenericType
print(generic_type.base)  # "List"
print(generic_type.parameters)  # ["str"]

# Modify parameters
generic_type.parameters.append("int")
generic_type.parameters[0] = "float"

# Create new generic
function.set_return_type("List[str]")
Learn more about working with collections here

​
Type Resolution

Type resolution uses Type.resolved_value to get the actual symbols that a type refers to:

# Get the actual symbols for a type
type_annotation = function.return_type  # -> Type
resolved_types = type_annotation.resolved_value  # Returns an Expression, likely a Symbol or collection of Symbols

# For generic types, resolve each parameter
if hasattr(type_annotation, "parameters"):
    for param in type_annotation.parameters:
        param_types = param.resolved_value # Get symbols for each parameter

# For union types, resolve each option
if hasattr(type_annotation, "options"):
    for option in type_annotation.options:
        option_types = option.resolved_value # Get symbols for each union option