|
| 1 | +--- |
| 2 | +title: "Working with Type Annotations" |
| 3 | +sidebarTitle: "Type Annotations" |
| 4 | +icon: "code" |
| 5 | +iconType: "solid" |
| 6 | +--- |
| 7 | + |
| 8 | +This guide covers the core APIs and patterns for working with type annotations in Codegen. |
| 9 | + |
| 10 | +## Type Resolution |
| 11 | + |
| 12 | +Codegen builds a complete dependency graph of your codebase, connecting functions, classes, imports, and their relationships. This enables powerful type resolution capabilities: |
| 13 | + |
| 14 | +```python |
| 15 | +from codegen import Codebase |
| 16 | + |
| 17 | +# Initialize codebase with dependency graph |
| 18 | +codebase = Codebase("./") |
| 19 | + |
| 20 | +# Get a function with a type annotation |
| 21 | +function = codebase.get_file("path/to/file.py").get_function("my_func") |
| 22 | + |
| 23 | +# Resolve its return type to actual symbols |
| 24 | +return_type = function.return_type |
| 25 | +resolved_symbols = return_type.resolved_types # Returns the actual Symbol objects |
| 26 | + |
| 27 | +# For generic types, you can resolve parameters |
| 28 | +if hasattr(return_type, "parameters"): |
| 29 | + for param in return_type.parameters: |
| 30 | + resolved_param = param.resolved_types # Get the actual type parameter symbols |
| 31 | + |
| 32 | +# For assignments, resolve their type |
| 33 | +assignment = codebase.get_file("path/to/file.py").get_assignment("my_var") |
| 34 | +resolved_type = assignment.type.resolved_types |
| 35 | +``` |
| 36 | + |
| 37 | +<Tip> |
| 38 | + Type resolution follows imports and handles complex cases like type aliases, forward references, and generic type parameters. |
| 39 | +</Tip> |
| 40 | + |
| 41 | +## Core Interfaces |
| 42 | + |
| 43 | +Type annotations in Codegen are built on two key interfaces: |
| 44 | + |
| 45 | +- [Typeable](/api-reference/core/Typeable) - The base interface for any node that can have a type annotation (parameters, variables, functions, etc). Provides `.type` and `.is_typed`. |
| 46 | +- [Type](/api-reference/core/Type) - The base class for all type annotations. Provides type resolution and dependency tracking. |
| 47 | + |
| 48 | +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. |
| 49 | + |
| 50 | +<Tip>Learn more about [inheritable behaviors](/building-with-codegen/inheritable-behaviors) like Typeable here</Tip> |
| 51 | + |
| 52 | +## Core Type APIs |
| 53 | + |
| 54 | +Type annotations can be accessed and modified through several key APIs: |
| 55 | + |
| 56 | +### Function Types |
| 57 | + |
| 58 | +The main APIs for function types are [Function.return_type](/api-reference/python/PyFunction#return-type) and [Function.set_return_type](/api-reference/python/PyFunction#set-return-type): |
| 59 | + |
| 60 | +```python |
| 61 | +# Get return type |
| 62 | +return_type = function.return_type # -> TypeAnnotation |
| 63 | +print(return_type.source) # "List[str]" |
| 64 | +print(return_type.is_typed) # True/False |
| 65 | + |
| 66 | +# Set return type |
| 67 | +function.set_return_type("List[str]") |
| 68 | +function.set_return_type(None) # Removes type annotation |
| 69 | +``` |
| 70 | + |
| 71 | +### Parameter Types |
| 72 | + |
| 73 | +Parameters use [Parameter.type](/api-reference/core/Parameter#type) and [Parameter.set_type_annotation](/api-reference/core/Parameter#set-type-annotation): |
| 74 | + |
| 75 | +```python |
| 76 | +for param in function.parameters: |
| 77 | + # Get parameter type |
| 78 | + param_type = param.type # -> TypeAnnotation |
| 79 | + print(param_type.source) # "int" |
| 80 | + print(param_type.is_typed) # True/False |
| 81 | + |
| 82 | + # Set parameter type |
| 83 | + param.set_type("int") |
| 84 | + param.set_type(None) # Removes type annotation |
| 85 | +``` |
| 86 | + |
| 87 | +### Variable Types |
| 88 | + |
| 89 | +Variables and attributes use [Assignment.type](/api-reference/core/Assignment#type) and [Assignment.set_type_annotation](/api-reference/core/Assignment#set-type-annotation). This applies to: |
| 90 | +- Global variables |
| 91 | +- Local variables |
| 92 | +- Class attributes (via [Class.attributes](/api-reference/core/Class#attributes)) |
| 93 | + |
| 94 | +```python |
| 95 | +# For global/local assignments |
| 96 | +assignment = file.get_assignment("my_var") |
| 97 | +var_type = assignment.type # -> TypeAnnotation |
| 98 | +print(var_type.source) # "str" |
| 99 | + |
| 100 | +# Set variable type |
| 101 | +assignment.set_type("str") |
| 102 | +assignment.set_type(None) # Removes type annotation |
| 103 | + |
| 104 | +# For class attributes |
| 105 | +class_def = file.get_class("MyClass") |
| 106 | +for attr in class_def.attributes: |
| 107 | + # Each attribute has an assignment property |
| 108 | + attr_type = attr.assignment.type # -> TypeAnnotation |
| 109 | + print(f"{attr.name}: {attr_type.source}") # e.g. "x: int" |
| 110 | + |
| 111 | + # Set attribute type |
| 112 | + attr.assignment.set_type("int") |
| 113 | + |
| 114 | +# You can also access attributes directly by index |
| 115 | +first_attr = class_def.attributes[0] |
| 116 | +first_attr.assignment.set_type("str") |
| 117 | +``` |
| 118 | + |
| 119 | +## Working with Complex Types |
| 120 | + |
| 121 | +### Union Types |
| 122 | + |
| 123 | +Union types ([UnionType](/api-reference/core/UnionType)) can be manipulated as collections: |
| 124 | + |
| 125 | +```python |
| 126 | +# Get union type |
| 127 | +union_type = function.return_type # -> A | B |
| 128 | +print(union_type.symbols) # ["A", "B"] |
| 129 | + |
| 130 | +# Add/remove options |
| 131 | +union_type.append("float") |
| 132 | +union_type.remove("None") |
| 133 | + |
| 134 | +# Check contents |
| 135 | +if "str" in union_type.options: |
| 136 | + print("String is a possible type") |
| 137 | +``` |
| 138 | +<Tip>Learn more about [working with collections here](/building-with-codegen/collections)</Tip> |
| 139 | + |
| 140 | +### Generic Types |
| 141 | + |
| 142 | +Generic types ([GenericType](/api-reference/core/GenericType)) expose their parameters as collection of [Parameters](/api-reference/core/Parameter): |
| 143 | + |
| 144 | +```python |
| 145 | +# Get generic type |
| 146 | +generic_type = function.return_type # -> GenericType |
| 147 | +print(generic_type.base) # "List" |
| 148 | +print(generic_type.parameters) # ["str"] |
| 149 | + |
| 150 | +# Modify parameters |
| 151 | +generic_type.parameters.append("int") |
| 152 | +generic_type.parameters[0] = "float" |
| 153 | + |
| 154 | +# Create new generic |
| 155 | +function.set_return_type("List[str]") |
| 156 | +``` |
| 157 | +<Tip>Learn more about [working with collections here](/building-with-codegen/collections)</Tip> |
| 158 | + |
| 159 | +### Type Resolution |
| 160 | + |
| 161 | +Type resolution uses [`Type.resolved_value`](/api-reference/core/Type#resolved-value) to get the actual symbols that a type refers to: |
| 162 | + |
| 163 | +```python |
| 164 | +# Get the actual symbols for a type |
| 165 | +type_annotation = function.return_type # -> Type |
| 166 | +resolved_types = type_annotation.resolved_value # Returns an Expression, likely a Symbol or collection of Symbols |
| 167 | + |
| 168 | +# For generic types, resolve each parameter |
| 169 | +if hasattr(type_annotation, "parameters"): |
| 170 | + for param in type_annotation.parameters: |
| 171 | + param_types = param.resolved_value # Get symbols for each parameter |
| 172 | + |
| 173 | +# For union types, resolve each option |
| 174 | +if hasattr(type_annotation, "options"): |
| 175 | + for option in type_annotation.options: |
| 176 | + option_types = option.resolved_value # Get symbols for each union option |
| 177 | +``` |
0 commit comments