From 40e2c8727f8cd1fea371edef32aa4887a37044c3 Mon Sep 17 00:00:00 2001 From: Shahin Saadati Date: Fri, 30 Jan 2026 13:32:04 -0800 Subject: [PATCH 1/2] Added proto text format output --- .gitignore | 2 + README.md | 10 +-- src/google/adk/scope/extractors/extract.py | 71 ++++++++++++++++++---- src/google/adk/scope/utils/args.py | 2 +- test/adk/scope/extractors/test_extract.py | 19 ++++-- test/adk/scope/utils/test_args.py | 4 +- 6 files changed, 83 insertions(+), 25 deletions(-) diff --git a/.gitignore b/.gitignore index b624be6..7da2e86 100644 --- a/.gitignore +++ b/.gitignore @@ -5,3 +5,5 @@ __pycache__/ *.pyc *.json + +output/ \ No newline at end of file diff --git a/README.md b/README.md index d866231..3ea906c 100644 --- a/README.md +++ b/README.md @@ -38,10 +38,10 @@ The script requires a `--language` argument to specify the target language (`py` ```bash # For Python -./extract.sh --language py --input-repo /path/to/adk-python output.json +./extract.sh --language py --input-repo /path/to/adk-python output_dir # For TypeScript -./extract.sh --language ts --input-repo /path/to/adk-js output.json +./extract.sh --language ts --input-repo /path/to/adk-js output_dir ``` ### CLI Arguments @@ -52,18 +52,18 @@ The script requires a `--language` argument to specify the target language (`py` | `--input-file ` | Path to a single file to process. | | `--input-dir ` | Path to a directory containing files. | | `--input-repo ` | Path to the root of an ADK repository. Recursive search in `src` (Python) or `core/src` (TS). | -| `output` | **Required.** Path to the output JSON file. | +| `output` | **Required.** Path to the output directory. | **Examples:** ```bash # Process a single file -./extract.sh --language python --input-file src/my_agent.py output.json +./extract.sh --language python --input-file src/my_agent.py output_dir # Process a directory python3 -m google.adk.scope.extractors.python.extractor \ --input-dir src/google/adk \ - output.json + output_dir ``` ## Development diff --git a/src/google/adk/scope/extractors/extract.py b/src/google/adk/scope/extractors/extract.py index 636f58f..cb338ee 100644 --- a/src/google/adk/scope/extractors/extract.py +++ b/src/google/adk/scope/extractors/extract.py @@ -2,7 +2,10 @@ import sys from pathlib import Path -from google.protobuf.json_format import MessageToJson +import datetime +import yaml +from google.protobuf import text_format +from google.protobuf.json_format import MessageToJson, MessageToDict from google.adk.scope.features_pb2 import FeatureRegistry from google.adk.scope.utils.args import parse_args from google.adk.scope.extractors import extractor_py, extractor_ts @@ -12,6 +15,12 @@ ) logger = logging.getLogger(__name__) +_JSON_INDENT = 2 +_JSON_OUTPUT = True +_YAML_OUTPUT = True +_PROTO_OUTPUT = True + + EXTRACTORS = { "python": extractor_py, "typescript": extractor_ts, @@ -161,19 +170,57 @@ def main(): features=all_features, ) + output_dir = args.output try: - with open(args.output, "w") as f: - f.write( - MessageToJson( - registry, - indent=2, - preserving_proto_field_name=True, - always_print_fields_with_no_presence=True, - ) - ) - logger.info("Successfully wrote output to %s", args.output) + output_dir.mkdir(parents=True, exist_ok=True) except IOError as e: - logger.error("Failed to write output: %s", e) + logger.error("Failed to create output directory %s: %s", output_dir, e) + sys.exit(1) + + timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") + prefix = "py" if args.language in {"python", "py"} else "ts" + base_filename = f"{prefix}_{timestamp}" + + if _JSON_OUTPUT: + # 1. JSON Output + json_path = output_dir / f"{base_filename}.json" + try: + json_content = MessageToJson( + registry, + indent=_JSON_INDENT, + preserving_proto_field_name=True, + always_print_fields_with_no_presence=True, + ) + json_path.write_text(json_content) + logger.info("Generated JSON: %s", json_path) + except IOError as e: + logger.error("Failed to write JSON output: %s", e) + + if _YAML_OUTPUT: + # 2. YAML Output + yaml_path = output_dir / f"{base_filename}.yaml" + try: + dict_content = MessageToDict( + registry, + preserving_proto_field_name=True, + always_print_fields_with_no_presence=True, + use_integers_for_enums=False, + ) + with open(yaml_path, "w") as f: + yaml.dump(dict_content, f, sort_keys=False) + logger.info("Generated YAML: %s", yaml_path) + except IOError as e: + logger.error("Failed to write YAML output: %s", e) + + if _PROTO_OUTPUT: + # 3. TextProto Output + txtpb_path = output_dir / f"{base_filename}.txtpb" + try: + txtpb_content = text_format.MessageToString(registry) + txtpb_path.write_text(txtpb_content) + logger.info("Generated TextProto: %s", txtpb_path) + except IOError as e: + logger.error("Failed to write TextProto output: %s", e) if __name__ == "__main__": diff --git a/src/google/adk/scope/utils/args.py b/src/google/adk/scope/utils/args.py index bf0bd7f..d544e0d 100644 --- a/src/google/adk/scope/utils/args.py +++ b/src/google/adk/scope/utils/args.py @@ -44,7 +44,7 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "output", type=Path, - help="Path to the output file", + help="Directory path for the output files", ) parser.add_argument( diff --git a/test/adk/scope/extractors/test_extract.py b/test/adk/scope/extractors/test_extract.py index dc8db06..d8dd013 100644 --- a/test/adk/scope/extractors/test_extract.py +++ b/test/adk/scope/extractors/test_extract.py @@ -67,7 +67,7 @@ class TestExtractMain(unittest.TestCase): def setUp(self): self.test_dir = tempfile.mkdtemp() self.root = Path(self.test_dir) - self.output_file = self.root / "output.json" + self.output_dir = self.root / "output" # Patches self.mock_args_patcher = patch( @@ -107,7 +107,7 @@ def configure_args( mock_args.input_file = Path(input_file) if input_file else None mock_args.input_dir = Path(input_dir) if input_dir else None mock_args.input_repo = Path(input_repo) if input_repo else None - mock_args.output = output or str(self.output_file) + mock_args.output = Path(output) if output else self.output_dir self.mock_args.return_value = mock_args def test_unsupported_language(self): @@ -133,7 +133,10 @@ def test_input_file_mode(self): extract.main() self.mock_py_extractor.extract_features.assert_called() - self.assertTrue(self.output_file.exists()) + self.assertTrue(self.output_dir.exists()) + self.assertEqual(len(list(self.output_dir.glob("py_*.json"))), 1) + self.assertEqual(len(list(self.output_dir.glob("py_*.yaml"))), 1) + self.assertEqual(len(list(self.output_dir.glob("py_*.txtpb"))), 1) def test_input_file_not_found(self): self.configure_args(lang="python", input_file="/non/existent.py") @@ -153,7 +156,10 @@ def test_input_dir_mode(self): extract.main() self.mock_py_extractor.find_files.assert_called_with(d, recursive=False) - self.assertTrue(self.output_file.exists()) + self.assertTrue(self.output_dir.exists()) + self.assertEqual(len(list(self.output_dir.glob("py_*.json"))), 1) + self.assertEqual(len(list(self.output_dir.glob("py_*.yaml"))), 1) + self.assertEqual(len(list(self.output_dir.glob("py_*.txtpb"))), 1) def test_input_repo_mode(self): r = self.root @@ -169,7 +175,10 @@ def test_input_repo_mode(self): self.mock_py_extractor.find_files.assert_called_with( r / "src", recursive=True ) - self.assertTrue(self.output_file.exists()) + self.assertTrue(self.output_dir.exists()) + self.assertEqual(len(list(self.output_dir.glob("py_*.json"))), 1) + self.assertEqual(len(list(self.output_dir.glob("py_*.yaml"))), 1) + self.assertEqual(len(list(self.output_dir.glob("py_*.txtpb"))), 1) if __name__ == "__main__": diff --git a/test/adk/scope/utils/test_args.py b/test/adk/scope/utils/test_args.py index dd8ef0e..0eef2a7 100644 --- a/test/adk/scope/utils/test_args.py +++ b/test/adk/scope/utils/test_args.py @@ -12,7 +12,7 @@ def test_parse_args(self, mock_parse): mock_args = argparse.Namespace( language="py", input_repo=Path("/tmp/repo"), - output=Path("/tmp/out.json"), + output=Path("/tmp/out_dir"), input_file=None, input_dir=None, ) @@ -24,7 +24,7 @@ def test_parse_args(self, mock_parse): args = parse_args() self.assertEqual(args.input_repo, Path("/tmp/repo")) - self.assertEqual(args.output, Path("/tmp/out.json")) + self.assertEqual(args.output, Path("/tmp/out_dir")) # Should be normalized self.assertEqual(args.language, "python") From dc929f7535b949f42de71d96ff3097731f6df229 Mon Sep 17 00:00:00 2001 From: Shahin Saadati Date: Fri, 30 Jan 2026 13:34:31 -0800 Subject: [PATCH 2/2] Added yaml to pyproject.toml --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index 4321663..8063fb5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -18,6 +18,7 @@ dependencies = [ "ruff", "pytest", "pytest-cov", + "pyyaml", ]