Working with Type Annotations

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

# 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

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

# 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")

# 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")

# 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]")

# 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