@arizeai/phoenix-mcp

import ast import sys from pathlib import Path from typing import Callable, Literal, Mapping, Sequence # ============================================================================= # String-to-DateTime field type conversions # ============================================================================= # Some fields in the schema are `str` but need to be converted to `datetime` # to reconcile with the data type that the client actually returns at runtime. # # This mapping defines which fields in which classes should have their type # annotations converted from string-based types to datetime during the # dataclass-to-TypedDict transformation process. # # Format: { # "ClassName": ["field_name1", "field_name2", ...], # ... # } # # Supported transformations: # - str → datetime # - Optional[str] → Optional[datetime] # - str = "default" → NotRequired[datetime] (with default removed) # - Optional[str] = None → NotRequired[datetime] (with Optional and default removed) # # Example: # Before: created_at: str # After: created_at: datetime # # Before: updated_at: Optional[str] = None # After: updated_at: NotRequired[datetime] STR_TO_DATETIME_ALTERATIONS = { "DatasetVersion": ["created_at"], } # ============================================================================= # AST Transformer to convert dataclass definitions to TypedDict definitions. # ============================================================================= class ConvertDataClassToTypedDict(ast.NodeTransformer): def __init__(self): self.current_class_name = None def visit_ImportFrom(self, node: ast.ImportFrom) -> ast.AST: """ Replace the dataclasses import with a TypedDict import from typing. """ if node.module == "dataclasses": return ast.ImportFrom( module="typing", names=[ast.alias(name="TypedDict", asname=None)], level=0, ) return node def visit_ClassDef(self, node: ast.ClassDef) -> ast.AST: """ Convert a class definition into a TypedDict definition. Also reorders the "type" field (if present) to the top for readability. """ # Remember the current class name before descending into fields old_class_name = self.current_class_name self.current_class_name = node.name # Visit and transform all statements in the class body. new_body = [self.visit(child) for child in node.body] # Look for a field named "type" defined as a Literal and move it to the front. for index, stmt in enumerate(new_body): if ( isinstance(stmt, ast.AnnAssign) and isinstance(stmt.target, ast.Name) and stmt.target.id == "type" and isinstance(stmt.annotation, ast.Subscript) and isinstance(stmt.annotation.value, ast.Name) and stmt.annotation.value.id == "Literal" and isinstance(stmt.annotation.slice, ast.Constant) ): new_body = [new_body[index]] + new_body[:index] + new_body[index + 1 :] break # Restore the previous class name self.current_class_name = old_class_name # Redefine the class so that it inherits from TypedDict. return ast.ClassDef( name=node.name, bases=[ast.Name(id="TypedDict", ctx=ast.Load())], keywords=node.keywords, body=new_body, decorator_list=[], ) def visit_AnnAssign(self, node: ast.AnnAssign) -> ast.AST: """ Process annotated assignments: - Convert str fields to datetime based on STR_TO_DATETIME_ALTERATIONS. - Rename fields ending with "_" (like schema_ or json_) by stripping the underscore. - Convert default values on fields (when present) to a NotRequired[...] annotation. - Change `type: str = "xyz"` into `type: Literal["xyz"]`. - If a field is Optional[...] with a default value, remove the Optional. """ # Convert str fields to datetime based on STR_TO_DATETIME_ALTERATIONS if ( self.current_class_name in STR_TO_DATETIME_ALTERATIONS and isinstance(node.target, ast.Name) and node.target.id in STR_TO_DATETIME_ALTERATIONS[self.current_class_name] ): # Handle direct str annotation if isinstance(node.annotation, ast.Name) and node.annotation.id == "str": node = ast.AnnAssign( target=node.target, annotation=ast.Name(id="datetime", ctx=ast.Load()), value=node.value, simple=node.simple, ) # Handle Optional[str] annotation elif ( isinstance(node.annotation, ast.Subscript) and isinstance(node.annotation.value, ast.Name) and node.annotation.value.id == "Optional" and isinstance(node.annotation.slice, ast.Name) and node.annotation.slice.id == "str" ): node = ast.AnnAssign( target=node.target, annotation=ast.Subscript( value=ast.Name(id="Optional", ctx=ast.Load()), slice=ast.Name(id="datetime", ctx=ast.Load()), ctx=ast.Load(), ), value=node.value, simple=node.simple, ) # Rename fields ending with "_" to remove the trailing underscore. if isinstance(node.target, ast.Name) and node.target.id in ("schema_", "json_"): new_target = ast.Name(id=node.target.id.rstrip("_"), ctx=ast.Store()) node = ast.AnnAssign( target=new_target, annotation=node.annotation, value=node.value, simple=node.simple, ) # If there is a default value, perform further transformations. if isinstance(node.value, ast.Constant): # Convert `type: str = "xyz"` into `type: Literal["xyz"]` if ( isinstance(node.target, ast.Name) and node.target.id == "type" and isinstance(node.annotation, ast.Name) and node.annotation.id == "str" ): return ast.AnnAssign( target=node.target, annotation=ast.Subscript( value=ast.Name(id="Literal", ctx=ast.Load()), slice=node.value, ctx=ast.Load(), ), value=None, # Remove default value simple=node.simple, ) # Convert an Optional annotation (with a default) to the inner type. if ( isinstance(node.annotation, ast.Subscript) and isinstance(node.annotation.value, ast.Name) and node.annotation.value.id == "Optional" ): node = ast.AnnAssign( target=node.target, annotation=node.annotation.slice, value=node.value, simple=node.simple, ) # Remove the default value and wrap the annotation in NotRequired. return ast.AnnAssign( target=node.target, annotation=ast.Subscript( value=ast.Name(id="NotRequired", ctx=ast.Load()), slice=node.annotation, ctx=ast.Load(), ), value=None, # Default value is removed simple=node.simple, ) return node def transform_dataclass(code: str) -> ast.AST: """ Parse the provided code, insert an import for NotRequired, and transform dataclass definitions into TypedDict definitions. Args: code: A string of Python source code. Returns: The transformed AST. """ parsed_ast: ast.Module = ast.parse(code) # Insert the imports for NotRequired and datetime before the first class. for index, node in enumerate(parsed_ast.body): if isinstance(node, ast.ClassDef): import_notrequired = ast.ImportFrom( module="typing_extensions", names=[ast.alias(name="NotRequired", asname=None)], level=0, ) import_datetime = ast.ImportFrom( module="datetime", names=[ast.alias(name="datetime", asname=None)], level=0, ) parsed_ast.body.insert(index, import_notrequired) parsed_ast.body.insert(index + 1, import_datetime) break # Remove top-level Union type definitions parsed_ast.body = [ node for node in parsed_ast.body if not ( isinstance(node, ast.Assign) and isinstance(node.value, ast.Subscript) and isinstance(node.value.value, ast.Name) and node.value.value.id == "Union" ) ] transformer = ConvertDataClassToTypedDict() transformed_ast = transformer.visit(parsed_ast) return transformed_ast # ============================================================================= # Functions to adjust class definitions by removing inherited fields. # ============================================================================= # Mapping from a class name to a list of its parent class names. PARENTS: Mapping[str, Sequence[str]] = { "Prompt": ["PromptData"], "PromptVersion": ["PromptVersionData"], "SpanAnnotation": ["SpanAnnotationData"], "LocalUser": ["LocalUserData"], "OAuth2User": ["OAuth2UserData"], } def get_ancestor_fields( class_name: str, class_nodes: Mapping[str, ast.ClassDef], parent_map: Mapping[str, Sequence[str]] = PARENTS, ) -> set[str]: """ Recursively collects the field names defined in all ancestor classes. Args: class_name: The name of the class to inspect. class_nodes: Mapping of class names to their AST ClassDef nodes. parent_map: Mapping from a class name to a list of its parent class names. Returns: A set of field names from all ancestors. """ if class_name not in parent_map: return set() fields: set[str] = set() for parent_name in parent_map[class_name]: parent_node: ast.ClassDef = class_nodes[parent_name] for stmt in parent_node.body: if isinstance(stmt, ast.AnnAssign) and isinstance(stmt.target, ast.Name): fields.add(stmt.target.id) # Recursively add fields from higher ancestors. fields |= get_ancestor_fields(parent_name, class_nodes, parent_map) return fields def remove_inherited_fields( class_nodes: Mapping[str, ast.ClassDef], parent_map: Mapping[str, Sequence[str]] = PARENTS, ) -> Mapping[str, ast.ClassDef]: """ For each class that inherits from others, remove any field that is already defined in its ancestors to improve readability. Args: class_nodes: Mapping from class name to its AST ClassDef node. parent_map: Mapping from class name to a list of its parent class names. Returns: A new mapping from class name to a modified AST ClassDef node. """ new_class_nodes: dict[str, ast.ClassDef] = {} for class_name, node in class_nodes.items(): # If the class has no parents, leave it as-is. if class_name not in parent_map: new_class_nodes[class_name] = node continue # Create explicit bases for the class from its parent names. bases: list[ast.expr] = [ ast.Name(id=parent, ctx=ast.Load()) for parent in parent_map[class_name] ] # Collect the field names defined in the class. child_field_names: set[str] = { stmt.target.id for stmt in node.body if isinstance(stmt, ast.AnnAssign) and isinstance(stmt.target, ast.Name) } # Ensure every statement in the body is an AnnAssign. assert len(child_field_names) == len(node.body), "Every field must be an AnnAssign" # Collect all ancestor field names. ancestor_field_names: set[str] = get_ancestor_fields(class_name, class_nodes, parent_map) assert ancestor_field_names < child_field_names, ( "Ancestor fields must be a subset of child fields" ) # Remove any inherited field from the class body. inherited_fields: set[str] = ancestor_field_names.intersection(child_field_names) new_body: list[ast.stmt] = [ stmt for stmt in node.body if not ( isinstance(stmt, ast.AnnAssign) and isinstance(stmt.target, ast.Name) and stmt.target.id in inherited_fields ) ] new_class_nodes[class_name] = ast.ClassDef( name=node.name, bases=bases, keywords=node.keywords, body=new_body, decorator_list=node.decorator_list, ) return new_class_nodes def topologically_sort_classes( class_nodes: Mapping[str, ast.ClassDef], parent_map: Mapping[str, Sequence[str]] = PARENTS, ) -> list[ast.ClassDef]: """ Sort class definitions so that each parent class appears before its children. Args: class_nodes: Mapping from class name to its AST ClassDef node. parent_map: Mapping from class name to a list of its parent class names. Returns: A list of AST ClassDef nodes in topological order. Raises: ValueError: If a cycle is detected in the inheritance hierarchy. """ sorted_classes: list[ast.ClassDef] = [] visit_state: dict[str, Literal["visiting", "visited"]] = {} def visit(node: ast.ClassDef) -> None: class_name: str = node.name if class_name in visit_state: if visit_state[class_name] == "visiting": raise ValueError(f"Cycle detected at class: {class_name}") return # Already visited visit_state[class_name] = "visiting" for parent_name in parent_map.get(class_name, []): visit(class_nodes[parent_name]) visit_state[class_name] = "visited" sorted_classes.append(node) for node in class_nodes.values(): visit(node) return sorted_classes # ============================================================================= # File rewriting logic. # ============================================================================= def rewrite_file( directory: Path, input_filename: str, output_filename: str, transform: Callable[[str], ast.AST], ) -> None: """ Reads a Python file, applies the AST transformation and class adjustments, then writes the updated code to an output file. Args: directory: The directory containing the input file. input_filename: The name of the input file. output_filename: The name of the output file. transform: A function that converts a code string to an AST. """ file_path: Path = directory / input_filename with open(file_path, "r") as file: code: str = file.read() transformed_ast: ast.AST = transform(code) assert isinstance(transformed_ast, ast.Module), "The transformed AST must be a module." # Extract all class definitions from the AST. class_nodes: dict[str, ast.ClassDef] = { node.name: node for node in transformed_ast.body if isinstance(node, ast.ClassDef) } # Remove inherited fields from subclasses. cleaned_classes: Mapping[str, ast.ClassDef] = remove_inherited_fields(class_nodes) # Topologically sort classes so that parent classes come first. sorted_classes: list[ast.ClassDef] = topologically_sort_classes(cleaned_classes) # Preserve non-class statements and then append the sorted class definitions. non_class_statements: list[ast.stmt] = [ stmt for stmt in transformed_ast.body if not isinstance(stmt, ast.ClassDef) ] new_body: list[ast.stmt] = non_class_statements + sorted_classes new_module: ast.Module = ast.Module(body=new_body, type_ignores=[]) new_module = ast.fix_missing_locations(new_module) output_code: str = ast.unparse(new_module) with open(directory / output_filename, "w") as file: file.write('"""Do not edit"""\n\n') file.write(output_code) # ============================================================================= # Main entry point. # ============================================================================= if __name__ == "__main__": if len(sys.argv) != 2: print("Usage: python transform.py <directory>") sys.exit(1) directory: Path = Path(sys.argv[1]) rewrite_file( directory, ".dataclass.py", "__init__.py", transform_dataclass, )