Skip to content

Commit d537285

Browse files
chore: Post processor script for grpc generation. (#170)
# Description Thank you for opening a Pull Request! Before submitting your PR, there are a few things you can do to make sure it goes smoothly: - [ ] Follow the [`CONTRIBUTING` Guide](https://github.com/google-a2a/a2a-python/blob/main/CONTRIBUTING.md). - [ ] Make your Pull Request title in the <https://www.conventionalcommits.org/> specification. - Important Prefixes for [release-please](https://github.com/googleapis/release-please): - `fix:` which represents bug fixes, and correlates to a [SemVer](https://semver.org/) patch. - `feat:` represents a new feature, and correlates to a SemVer minor. - `feat!:`, or `fix!:`, `refactor!:`, etc., which represent a breaking change (indicated by the `!`) and will result in a SemVer major. - [ ] Ensure the tests and linter pass (Run `nox -s format` from the repository root to format) - [ ] Appropriate docs were updated (if necessary) Fixes #<issue_number_goes_here> 🦕
1 parent e744e6d commit d537285

File tree

1 file changed

+56
-0
lines changed

1 file changed

+56
-0
lines changed

scripts/grpc_gen_post_processor.py

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
"""
2+
Fix absolute imports in *_pb2_grpc.py files.
3+
Example:
4+
import a2a_pb2 as a2a__pb2
5+
from . import a2a_pb2 as a2a__pb2
6+
"""
7+
8+
import re
9+
import sys
10+
11+
from pathlib import Path
12+
13+
14+
def process_generated_code(src_folder: str = 'src/a2a/grpc'):
15+
"""Post processor for the generated code."""
16+
dir_path = Path(src_folder)
17+
print(dir_path)
18+
if not dir_path.is_dir():
19+
print('Source folder not found')
20+
sys.exit(1)
21+
22+
grpc_pattern = '**/*_pb2_grpc.py'
23+
files = dir_path.glob(grpc_pattern)
24+
25+
for file in files:
26+
print(f'Processing {file}')
27+
try:
28+
with file.open('r', encoding='utf-8') as f:
29+
src_content = f.read()
30+
31+
# Change import a2a_pb2 as a2a__pb2
32+
import_pattern = r'^import (\w+_pb2) as (\w+__pb2)$'
33+
# to from . import a2a_pb2 as a2a__pb2
34+
replacement_pattern = r'from . import \1 as \2'
35+
36+
fixed_src_content = re.sub(
37+
import_pattern,
38+
replacement_pattern,
39+
src_content,
40+
flags=re.MULTILINE,
41+
)
42+
43+
if fixed_src_content != src_content:
44+
with file.open('w', encoding='utf-8') as f:
45+
f.write(fixed_src_content)
46+
print('Imports fixed')
47+
else:
48+
print('No changes needed')
49+
50+
except Exception as e:
51+
print(f'Error processing file {file}: {e}')
52+
sys.exit(1)
53+
54+
55+
if __name__ == '__main__':
56+
process_generated_code()

0 commit comments

Comments
 (0)