MCP Terminal Server
by dillip285
- py
- bin
#!/usr/bin/env python3
#
# Copyright 2025 Google LLC
# SPDX-License-Identifier: Apache-2.0
"""Standalone convenience script used to massage the typing.py.
The `py/packages/genkit/src/genkit/core/schema_types.py` file is generated by
datamodel-codegen. However, since the tool doesn't currently provide options to
generate exactly the kind of code we need, we use this convenience script to
parse the Python source code, walk the AST, modify it to include the bits we
need and regenerate the code for eventual use within our codebase.
Transformations applied:
- We remove the model_config attribute from classes that ineherit from
RootModel.
- We add the `populate_by_name=True` parameter to ensure serialization uses
camelCase for attributes since the JS implementation uses camelCase and Python
uses snake_case. The codegen pass is configured to generate snake_case for a
Pythonic API but serialize to camelCase in order to be compatible with
runtimes.
- We add a license header
- We add a header indicating that this file has been generated by a code
generator pass.
- We add the ability to use forward references.
- Add docstrings if missing.
"""
import ast
import sys
from datetime import datetime
from pathlib import Path
class ModelConfigRemover(ast.NodeTransformer):
"""AST transformer to remove model_config from RootModel classes.
This class traverses the AST and removes model_config assignments from
RootModel classes while preserving them in other classes.
"""
def __init__(self) -> None:
"""Initialize the ModelConfigRemover."""
self.modified = False
def is_rootmodel_class(self, node: ast.ClassDef) -> bool:
"""Check if a class definition is a RootModel class.
Args:
node: The AST node representing a class definition.
Returns:
True if the class inherits from RootModel, False otherwise.
"""
for base in node.bases:
if isinstance(base, ast.Name) and base.id == 'RootModel':
return True
elif isinstance(base, ast.Subscript):
value = base.value
if isinstance(value, ast.Name) and value.id == 'RootModel':
return True
return False
def create_model_config(
self, extra_forbid: bool = True, populate_by_name: bool = True
) -> ast.Assign:
"""Create a model_config assignment with the specified options."""
keywords = []
if extra_forbid:
keywords.append(
ast.keyword(arg='extra', value=ast.Constant(value='forbid'))
)
if populate_by_name:
keywords.append(
ast.keyword(
arg='populate_by_name', value=ast.Constant(value=True)
)
)
return ast.Assign(
targets=[ast.Name(id='model_config')],
value=ast.Call(
func=ast.Name(id='ConfigDict'), args=[], keywords=keywords
),
)
def has_model_config(self, node: ast.ClassDef) -> bool:
"""Check if a class already has a model_config assignment."""
for item in node.body:
if isinstance(item, ast.Assign):
targets = item.targets
if len(targets) == 1 and isinstance(targets[0], ast.Name):
if targets[0].id == 'model_config':
return True
return False
def visit_ClassDef(self, node: ast.ClassDef) -> ast.ClassDef: # noqa: N802
"""Visit class definitions and handle model_config based on class type.
This method removes model_config assignments from RootModel classes
while preserving them in other classes.
Args:
node: The AST node representing a class definition.
Returns:
The modified class definition node.
"""
if self.is_rootmodel_class(node):
# Filter out model_config assignments for RootModel classes
node.body = [
stmt
for stmt in node.body
if not (
isinstance(stmt, ast.Assign)
and any(
isinstance(target, ast.Name)
and target.id == 'model_config'
for target in stmt.targets
)
)
]
self.modified = True
else:
# For non-RootModel classes that inherit from BaseModel
if any(
isinstance(base, ast.Name) and base.id == 'BaseModel'
for base in node.bases
):
if not self.has_model_config(node):
# Add model_config with populate_by_name=True
node.body.insert(0, self.create_model_config())
self.modified = True
else:
# Update existing model_config to include
# populate_by_name=True
new_body = []
for item in node.body:
if isinstance(item, ast.Assign):
targets = item.targets
if len(targets) == 1 and isinstance(
targets[0], ast.Name
):
if targets[0].id == 'model_config':
new_body.append(self.create_model_config())
self.modified = True
continue
new_body.append(item)
node.body = new_body
return node
class ClassTransformer(ast.NodeTransformer):
"""AST transformer that modifies class definitions."""
def visit_ClassDef(self, node: ast.ClassDef) -> ast.ClassDef: # noqa: N802
"""Visit and transform a class definition node.
Args:
node: The ClassDef AST node to transform.
Returns:
The transformed ClassDef node.
"""
# First apply base class transformations
node = super().generic_visit(node)
# Check if class has a docstring
has_docstring = False
if node.body and isinstance(node.body[0], ast.Expr):
if isinstance(node.body[0].value, ast.Str):
has_docstring = True
# Add docstring if missing
if not has_docstring:
# Generate a more descriptive docstring based on class type
if any(
base.id == 'RootModel'
for base in node.bases
if isinstance(base, ast.Name)
):
docstring = (
f'Root model for {node.name.lower().replace("_", " ")}.'
)
elif any(
base.id == 'BaseModel'
for base in node.bases
if isinstance(base, ast.Name)
):
docstring = (
f'Model for {node.name.lower().replace("_", " ")} data.'
)
elif any(
base.id == 'Enum'
for base in node.bases
if isinstance(base, ast.Name)
):
n = node.name.lower().replace('_', ' ')
docstring = f'Enumeration of {n} values.'
else:
docstring = f'{node.name} data type class.'
node.body.insert(0, ast.Expr(value=ast.Str(s=docstring)))
# Remove model_config from RootModel classes
if any(
base.id == 'RootModel'
for base in node.bases
if isinstance(base, ast.Name)
):
node.body = [
stmt
for stmt in node.body
if not (
isinstance(stmt, ast.AnnAssign)
and isinstance(stmt.target, ast.Name)
and stmt.target.id == 'model_config'
)
]
return node
def add_header(content: str) -> str:
"""Add the generated header to the content."""
header = '''# Copyright {year} Google LLC
# SPDX-License-Identifier: Apache-2.0
#
# DO NOT EDIT: Generated by `generate_schema_typing` from `genkit-schemas.json`.
"""Schema types module defining the core data models for Genkit.
This module contains Pydantic models that define the structure and validation
for various data types used throughout the Genkit framework, including messages,
actions, tools, and configuration options.
"""
from __future__ import annotations
'''
return header.format(year=datetime.now().year) + content
def process_file(filename: str) -> None:
"""Process a Python file to remove model_config from RootModel classes.
This function reads a Python file, processes its AST to remove model_config
from RootModel classes, and writes the modified code back to the file.
Args:
filename: Path to the Python file to process.
Raises:
FileNotFoundError: If the input file does not exist.
SyntaxError: If the input file contains invalid Python syntax.
"""
path = Path(filename)
if not path.is_file():
print(f'Error: File not found: {filename}')
sys.exit(1)
try:
with open(path, encoding='utf-8') as f:
source = f.read()
tree = ast.parse(source)
model_transformer = ModelConfigRemover()
modified_tree = model_transformer.visit(tree)
class_transformer = ClassTransformer()
modified_tree = class_transformer.visit(modified_tree)
if (
hasattr(model_transformer, 'modified')
and model_transformer.modified
) or (
hasattr(class_transformer, 'modified')
and class_transformer.modified
):
# Format the modified code
ast.fix_missing_locations(modified_tree)
modified_source = ast.unparse(modified_tree)
src = add_header(modified_source)
with open(path, 'w', encoding='utf-8') as f:
f.write(src)
print(f'Successfully processed {filename}')
else:
# Even if no AST modifications, still write back to add the header.
path.write_text(add_header(source), encoding='utf-8')
print(f'Added header to {filename}')
except SyntaxError as e:
print(f'Error: Invalid Python syntax in {filename}: {e}')
sys.exit(1)
def main() -> None:
"""Main entry point for the script.
This function processes command line arguments and calls the appropriate
functions to process the schema types file.
Usage:
python script.py <filename>
"""
if len(sys.argv) != 2:
print('Usage: python script.py <filename>')
sys.exit(1)
process_file(sys.argv[1])
if __name__ == '__main__':
main()