Implementation:Arize ai Phoenix Codegen Transform

Knowledge Sources	Arize_ai_Phoenix
Domains	AI_Observability, Client_SDK, Code_Generation
Last Updated	2026-02-14 05:30 GMT

Overview

AST-based code generation script that converts Python dataclass definitions into TypedDict definitions, producing the auto-generated V1 API types for the Phoenix client.

Description

The Codegen Transform module is a build-time tool that transforms a source file of Python dataclass definitions (generated from the Phoenix OpenAPI specification) into a file of TypedDict definitions suitable for use as a type-safe client-side API schema.

The module contains three major components:

1. ConvertDataClassToTypedDict (AST Transformer):

This ast.NodeTransformer subclass performs the core transformation by visiting and rewriting AST nodes:

• visit_ImportFrom() -- Replaces from dataclasses import ... with from typing import TypedDict.
• visit_ClassDef() -- Rewrites each class to inherit from TypedDict instead of dataclass. Reorders the type field (if defined as a Literal) to the top of the class body for readability.
• visit_AnnAssign() -- Performs several field-level transformations:
  • str-to-datetime conversion: For fields listed in STR_TO_DATETIME_ALTERATIONS (e.g., DatasetVersion.created_at), rewrites str annotations to datetime and Optional[str] to Optional[datetime].
  • Field renaming: Strips trailing underscores from fields like schema_ and json_ (Python naming workaround for reserved words).
  • Literal conversion: Transforms type: str = "xyz" into type: Literal["xyz"].
  • NotRequired wrapping: For any field with a default value, removes the default and wraps the annotation in NotRequired[...].
  • Optional unwrapping: For Optional[T] fields with defaults, extracts the inner type before wrapping in NotRequired.

2. Inheritance Processing Functions:

• get_ancestor_fields() -- Recursively collects field names from all ancestor classes based on a predefined parent mapping (PARENTS dict).
• remove_inherited_fields() -- For classes with parents (e.g., Prompt extends PromptData), removes fields already defined in ancestors to avoid duplication in the TypedDict output.
• topologically_sort_classes() -- Ensures parent classes are defined before their children in the output file, using a DFS-based topological sort with cycle detection.

3. File Rewriting:

• rewrite_file() -- Orchestrates the full pipeline: reads the input file (.dataclass.py), applies the AST transformation, removes inherited fields, topologically sorts classes, and writes the output (__init__.py) with a "Do not edit" header.
• transform_dataclass() -- Parses the source code, injects imports for NotRequired and datetime, removes top-level Union type aliases, and runs the AST transformer.

The PARENTS mapping defines the inheritance hierarchy for classes that use TypedDict inheritance (e.g., Prompt -> PromptData, SpanAnnotation -> SpanAnnotationData, LocalUser -> LocalUserData).

Usage

Run this script during the build process to regenerate the V1 types after modifying the OpenAPI specification or the intermediate dataclass definitions. It is invoked as python transform.py <directory> where the directory contains the .dataclass.py input file.

Code Reference

Source Location

• Repository: Arize_ai_Phoenix
• File: packages/phoenix-client/scripts/codegen/transform.py

Signature

class ConvertDataClassToTypedDict(ast.NodeTransformer):
    def __init__(self): ...
    def visit_ImportFrom(self, node: ast.ImportFrom) -> ast.AST: ...
    def visit_ClassDef(self, node: ast.ClassDef) -> ast.AST: ...
    def visit_AnnAssign(self, node: ast.AnnAssign) -> ast.AST: ...

def transform_dataclass(code: str) -> ast.AST: ...

def get_ancestor_fields(
    class_name: str,
    class_nodes: Mapping[str, ast.ClassDef],
    parent_map: Mapping[str, Sequence[str]] = PARENTS,
) -> set[str]: ...

def remove_inherited_fields(
    class_nodes: Mapping[str, ast.ClassDef],
    parent_map: Mapping[str, Sequence[str]] = PARENTS,
) -> Mapping[str, ast.ClassDef]: ...

def topologically_sort_classes(
    class_nodes: Mapping[str, ast.ClassDef],
    parent_map: Mapping[str, Sequence[str]] = PARENTS,
) -> list[ast.ClassDef]: ...

def rewrite_file(
    directory: Path,
    input_filename: str,
    output_filename: str,
    transform: Callable[[str], ast.AST],
) -> None: ...

Import

# This is a build script, not a runtime module.
# Invoked from the command line:
# python packages/phoenix-client/scripts/codegen/transform.py <directory>

I/O Contract

Input

Parameter	Type	Description
directory	Path	Directory containing the .dataclass.py source file
Input file	.dataclass.py	Python file with @dataclass definitions generated from OpenAPI spec

Output

Output	Type	Description
Output file	__init__.py	Python file with TypedDict definitions, prefixed with "Do not edit"

Transformation Rules

Input Pattern	Output Pattern	Description
@dataclass class Foo:	class Foo(TypedDict):	Dataclass to TypedDict conversion
from dataclasses import ...	from typing import TypedDict	Import rewriting
field: str = "default"	field: NotRequired[str]	Default values become NotRequired
type: str = "xyz"	type: Literal["xyz"]	Discriminator fields become Literals
field: Optional[T] = None	field: NotRequired[T]	Optional with default becomes NotRequired
created_at: str (in DatasetVersion)	created_at: datetime	Configured str-to-datetime conversion
schema_: str	schema: str	Trailing underscore removed from reserved-word fields

Parent Inheritance Map

Child Class	Parent Class(es)	Effect
Prompt	PromptData	Inherits name, description, source_prompt_id, metadata
PromptVersion	PromptVersionData	Inherits all version data fields
SpanAnnotation	SpanAnnotationData	Inherits name, annotator_kind, span_id, result, metadata, identifier
LocalUser	LocalUserData	Inherits email, username, role, auth_method, password
OAuth2User	OAuth2UserData	Inherits email, username, role, auth_method, oauth2 fields

Usage Examples

# Command-line invocation
# python packages/phoenix-client/scripts/codegen/transform.py \
#     packages/phoenix-client/src/phoenix/client/__generated__/v1/

# Programmatic usage (for testing or custom pipelines)
from pathlib import Path

# This would be in a build script context
directory = Path("packages/phoenix-client/src/phoenix/client/__generated__/v1/")

# The script reads .dataclass.py and writes __init__.py
# Input:
#   @dataclass
#   class Dataset:
#       id: str
#       name: str
#       description: Optional[str] = None
#
# Output:
#   class Dataset(TypedDict):
#       id: str
#       name: str
#       description: NotRequired[str]

Related Pages

• Principle:Arize_ai_Phoenix_Client_Initialization
• Arize_ai_Phoenix_Generated_V1_Types -- The output produced by this transform script