Skip to content

Commit 646b0a7

Browse files
Jacksunweicopybara-github
authored andcommitted
chore(config): Creates yaml_utils.py in utils to allow adk dump yaml in the same style
PiperOrigin-RevId: 795628807
1 parent f03f167 commit 646b0a7

File tree

2 files changed

+145
-0
lines changed

2 files changed

+145
-0
lines changed

src/google/adk/utils/yaml_utils.py

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
from __future__ import annotations
16+
17+
from pathlib import Path
18+
from typing import Union
19+
20+
from pydantic import BaseModel
21+
import yaml
22+
23+
24+
def dump_pydantic_to_yaml(
25+
model: BaseModel,
26+
file_path: Union[str, Path],
27+
*,
28+
indent: int = 2,
29+
sort_keys: bool = True,
30+
exclude_none: bool = True,
31+
) -> None:
32+
"""Dump a Pydantic model to a YAML file with multiline strings using | style.
33+
34+
Args:
35+
model: The Pydantic model instance to dump.
36+
file_path: Path to the output YAML file.
37+
indent: Number of spaces for indentation (default: 2).
38+
sort_keys: Whether to sort dictionary keys (default: True).
39+
exclude_none: Exclude fields with None values (default: True).
40+
"""
41+
model_dict = model.model_dump(exclude_none=exclude_none, mode='json')
42+
43+
file_path = Path(file_path)
44+
file_path.parent.mkdir(parents=True, exist_ok=True)
45+
46+
# Create a custom dumper class
47+
class _MultilineDumper(yaml.SafeDumper):
48+
pass
49+
50+
def multiline_str_representer(dumper, data):
51+
if '\n' in data:
52+
return dumper.represent_scalar('tag:yaml.org,2002:str', data, style='|')
53+
return dumper.represent_scalar('tag:yaml.org,2002:str', data)
54+
55+
# Add representer only to our custom dumper
56+
_MultilineDumper.add_representer(str, multiline_str_representer)
57+
58+
with file_path.open('w', encoding='utf-8') as f:
59+
yaml.dump(
60+
model_dict,
61+
f,
62+
Dumper=_MultilineDumper,
63+
indent=indent,
64+
sort_keys=sort_keys,
65+
default_flow_style=False,
66+
)
Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,79 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
"""Tests for YAML utility functions."""
16+
17+
from pathlib import Path
18+
from typing import Optional
19+
20+
from google.adk.utils.yaml_utils import dump_pydantic_to_yaml
21+
from google.genai import types
22+
from pydantic import BaseModel
23+
24+
25+
class SimpleModel(BaseModel):
26+
"""Simple test model."""
27+
28+
name: str
29+
age: int
30+
active: bool
31+
finish_reason: Optional[types.FinishReason] = None
32+
multiline_text: Optional[str] = None
33+
34+
35+
def test_yaml_file_generation(tmp_path: Path):
36+
"""Test that YAML file is correctly generated."""
37+
model = SimpleModel(
38+
name="Alice",
39+
age=30,
40+
active=True,
41+
finish_reason=types.FinishReason.STOP,
42+
)
43+
yaml_file = tmp_path / "test.yaml"
44+
45+
dump_pydantic_to_yaml(model, yaml_file)
46+
47+
assert yaml_file.read_text(encoding="utf-8") == """\
48+
active: true
49+
age: 30
50+
finish_reason: STOP
51+
name: Alice
52+
"""
53+
54+
55+
def test_multiline_string_pipe_style(tmp_path: Path):
56+
"""Test that multiline strings use | style."""
57+
multiline_text = """\
58+
This is a long description
59+
that spans multiple lines
60+
and should be formatted with pipe style"""
61+
model = SimpleModel(
62+
name="Test",
63+
age=25,
64+
active=False,
65+
multiline_text=multiline_text,
66+
)
67+
yaml_file = tmp_path / "test.yaml"
68+
69+
dump_pydantic_to_yaml(model, yaml_file)
70+
71+
assert yaml_file.read_text(encoding="utf-8") == """\
72+
active: false
73+
age: 25
74+
multiline_text: |-
75+
This is a long description
76+
that spans multiple lines
77+
and should be formatted with pipe style
78+
name: Test
79+
"""

0 commit comments

Comments
 (0)