-
Notifications
You must be signed in to change notification settings - Fork 151
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat(py): generate
schemas.py
from genkit-schema.json
using datam…
…odel-codegen #1807 (#1808) ISSUE: #1807 CHANGELOG: - [ ] Add a `bin/generate_schema_types` script that generates the Pydantic `schemas.py` module. - [ ] Not updating the pre-commit hooks since the generated file requires manual patching. - [ ] Remove timestamp to ensure we do not treat a file with identical content differently preventing the hassles of updating this file per commit.
- Loading branch information
Showing
10 changed files
with
566 additions
and
216 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,3 @@ | ||
# Copyright 2025 Google LLC | ||
# SPDX-License-Identifier: Apache-2.0 | ||
|
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,35 @@ | ||
#!/usr/bin/env bash | ||
# | ||
# Copyright 2025 Google LLC | ||
# SPDX-License-Identifier: Apache-2.0 | ||
|
||
set -euo pipefail | ||
|
||
TOP_DIR=$(git rev-parse --show-toplevel) | ||
SCHEMA_FILE="$TOP_DIR/py/packages/genkit/src/genkit/core/schemas.py" | ||
|
||
# Generate types using configuration from pyproject.toml | ||
uv run --directory "$TOP_DIR/py" datamodel-codegen | ||
|
||
# This isn't causing runtime errors at the moment so letting it be. | ||
#sed -i '' '/^class Model(RootModel\[Any\]):$/,/^ root: Any$/d' "$SCHEMA_FILE" | ||
|
||
# Sanitize the generated schema. | ||
python3 "${TOP_DIR}/py/bin/sanitize_schemas.py" "$SCHEMA_FILE" | ||
|
||
# Add a generated by `generate_schema_types` comment. | ||
sed -i '' '1i\ | ||
# DO NOT EDIT: Generated by `generate_schema_types` from `genkit-schemas.json`. | ||
' "$SCHEMA_FILE" | ||
|
||
# Add license header. | ||
addlicense \ | ||
-c "Google LLC" \ | ||
-s=only \ | ||
"$SCHEMA_FILE" | ||
|
||
# Checks and formatting. | ||
uv run --directory "$TOP_DIR/py" \ | ||
ruff check --fix "$SCHEMA_FILE" | ||
uv run --directory "$TOP_DIR/py" \ | ||
ruff format "$SCHEMA_FILE" |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,85 @@ | ||
#!/usr/bin/env python3 | ||
# | ||
# Copyright 2025 Google LLC | ||
# SPDX-License-Identifier: Apache-2.0 | ||
|
||
import ast | ||
import sys | ||
|
||
|
||
class ModelConfigRemover(ast.NodeTransformer): | ||
def __init__(self) -> None: | ||
self.modified = False | ||
|
||
def is_rootmodel_class(self, node: ast.ClassDef) -> bool: | ||
"""Check if the class inherits from RootModel.""" | ||
for base in node.bases: | ||
if isinstance(base, ast.Name) and base.id == 'RootModel': | ||
return True | ||
elif isinstance(base, ast.Subscript): | ||
if ( | ||
isinstance(base.value, ast.Name) | ||
and base.value.id == 'RootModel' | ||
): | ||
return True | ||
return False | ||
|
||
def visit_ClassDef(self, node: ast.ClassDef) -> ast.ClassDef: | ||
"""Visit class definitions and remove model_config if class inherits from RootModel.""" | ||
if self.is_rootmodel_class(node): | ||
# Filter out model_config assignments | ||
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(item) | ||
else: | ||
new_body.append(item) | ||
|
||
if len(new_body) != len(node.body): | ||
self.modified = True | ||
|
||
node.body = new_body | ||
|
||
return node | ||
|
||
|
||
def process_file(filename: str) -> None: | ||
"""Process a Python file to remove model_config from RootModel classes.""" | ||
with open(filename, 'r') as f: | ||
source = f.read() | ||
|
||
tree = ast.parse(source) | ||
|
||
transformer = ModelConfigRemover() | ||
modified_tree = transformer.visit(tree) | ||
|
||
if transformer.modified: | ||
ast.fix_missing_locations(modified_tree) | ||
modified_source = ast.unparse(modified_tree) | ||
with open(filename, 'w') as f: | ||
f.write(modified_source) | ||
print( | ||
f'Modified {filename}: Removed model_config from RootModel classes' | ||
) | ||
else: | ||
print(f'No modifications needed in {filename}') | ||
|
||
|
||
def main() -> None: | ||
if len(sys.argv) != 2: | ||
print('Usage: python script.py <filename>') | ||
sys.exit(1) | ||
|
||
filename = sys.argv[1] | ||
try: | ||
process_file(filename) | ||
except Exception as e: | ||
print(f'Error processing {filename}: {str(e)}') | ||
sys.exit(1) | ||
|
||
|
||
if __name__ == '__main__': | ||
main() |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.