Working with Type Annotations

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:

• Global variables
• Local variables
• Class attributes (via Class.attributes)

# 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