Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 1 addition & 2 deletions docs/building-with-codegen/calling-out-to-llms.mdx
Original file line number Diff line number Diff line change
Expand Up @@ -225,8 +225,7 @@ method.edit(new_impl)
```python
codebase.set_session_options(max_ai_requests=200)
```
<
Note>
<Note>
You can also use `codebase.set_session_options` to increase the execution time and the number of operations allowed in a session. This is useful for handling larger tasks or more complex operations that require additional resources. Adjust the `max_seconds` and `max_transactions` parameters to suit your needs:
```python
codebase.set_session_options(max_seconds=300, max_transactions=500)
Expand Down
2 changes: 1 addition & 1 deletion docs/building-with-codegen/reducing-conditions.mdx
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
---
title: "Reducing Conditions (Advanced)"
title: "Reducing Conditions"
sidebarTitle: "Reducing Conditions"
icon: "code-branch"
iconType: "solid"
Expand Down
177 changes: 177 additions & 0 deletions docs/building-with-codegen/type-annotations.mdx
Original file line number Diff line number Diff line change
@@ -0,0 +1,177 @@
---
title: "Working with Type Annotations"
sidebarTitle: "Type Annotations"
icon: "code"
iconType: "solid"
---

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:

```python
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
```

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

## Core Interfaces

Type annotations in Codegen are built on two key interfaces:

- [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`.
- [Type](/api-reference/core/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.

<Tip>Learn more about [inheritable behaviors](/building-with-codegen/inheritable-behaviors) like Typeable here</Tip>

## 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](/api-reference/python/PyFunction#return-type) and [Function.set_return_type](/api-reference/python/PyFunction#set-return-type):

```python
# 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](/api-reference/core/Parameter#type) and [Parameter.set_type_annotation](/api-reference/core/Parameter#set-type-annotation):

```python
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](/api-reference/core/Assignment#type) and [Assignment.set_type_annotation](/api-reference/core/Assignment#set-type-annotation). This applies to:
- Global variables
- Local variables
- Class attributes (via [Class.attributes](/api-reference/core/Class#attributes))

```python
# 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](/api-reference/core/UnionType)) can be manipulated as collections:

```python
# 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")
```
<Tip>Learn more about [working with collections here](/building-with-codegen/collections)</Tip>

### Generic Types

Generic types ([GenericType](/api-reference/core/GenericType)) expose their parameters as collection of [Parameters](/api-reference/core/Parameter):

```python
# 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]")
```
<Tip>Learn more about [working with collections here](/building-with-codegen/collections)</Tip>

### Type Resolution

Type resolution uses [`Type.resolved_value`](/api-reference/core/Type#resolved-value) to get the actual symbols that a type refers to:

```python
# 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
```
3 changes: 3 additions & 0 deletions docs/introduction/faq.mdx
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,9 @@ iconType: "solid"
Interested in adding support for your language? [Let us know](https://x.com/codegen) or [contribute](/introduction/community)!
</Note>

</Accordion>
<Accordion title="Is Codegen exact?" icon="scale-balanced">
Pretty much! Codegen is roughly on par with `mypy` and `tsc`. There are always edge cases in static analysis that are provably impossible to get (for example doing `eval()` on a string), but all of Codegen's APIs are intended to be exact unless otherwise specified. Please reach out if you find an edge case and we will do our best to patch it.
</Accordion>
<Accordion title="Is Codegen suitable for large codebases?" icon="database">
Yes! Codegen was developed on multmillion-line Python and Typescript codebases
Expand Down
10 changes: 6 additions & 4 deletions docs/introduction/overview.mdx
Original file line number Diff line number Diff line change
Expand Up @@ -9,9 +9,11 @@ iconType: "solid"

It provides a scriptable interface to a powerful, multi-lingual language server built on top of [Tree-sitter](https://tree-sitter.github.io/tree-sitter/).

export const metaCode = `# Codegen builds a complete graph connecting
# functions, classes, imports and their relationships
from codegen import Codebase
export const metaCode = `from codegen import Codebase

# Codegen builds a complete graph connecting
# functions, classes, imports and their relationships
codebase = Codebase("./")

# Work with code without dealing with syntax trees or parsing
for function in codebase.functions:
Expand All @@ -33,7 +35,7 @@ def baz():

<iframe
width="100%"
height="300px"
height="320px"
scrolling="no"
src={`https://chadcode.sh/embedded/codemod/?code=${encodeURIComponent(
metaCode
Expand Down
1 change: 1 addition & 0 deletions docs/mint.json
Original file line number Diff line number Diff line change
Expand Up @@ -111,6 +111,7 @@
"building-with-codegen/local-variables",
"building-with-codegen/comments-and-docstrings",
"building-with-codegen/external-modules",
"building-with-codegen/type-annotations",
"building-with-codegen/moving-symbols",
"building-with-codegen/collections",
"building-with-codegen/traversing-the-call-graph",
Expand Down
Loading
Loading