5151TRACE_LEVEL = 5
5252logging .addLevelName (TRACE_LEVEL , "TRACE" )
5353
54+
5455def trace (msg , * args , ** kwargs ):
5556 logging .log (TRACE_LEVEL , msg , * args , ** kwargs )
5657
58+
5759logging .trace = trace
5860
61+
5962@dataclass (frozen = True )
6063class Config :
6164 hapi_version : str
@@ -64,6 +67,7 @@ class Config:
6467 mirror_out : Path
6568 pyi_out : bool = True
6669
70+
6771def setup_logging (verbosity : int ) -> None :
6872 level = logging .WARNING
6973 if verbosity == 1 :
@@ -74,11 +78,14 @@ def setup_logging(verbosity: int) -> None:
7478 level = TRACE_LEVEL
7579 logging .basicConfig (level = level , format = "%(levelname)s: %(message)s" )
7680
81+
7782def resolve_path (p : str ) -> Path :
7883 q = Path (p )
7984 return q if q .is_absolute () else (SCRIPT_DIR / q )
8085
8186# -------------------- CLI --------------------
87+
88+
8289def parse_args () -> argparse .Namespace :
8390 parser = argparse .ArgumentParser (
8491 description = "Generate Python protobuf files from Hedera proto definitions."
@@ -117,6 +124,8 @@ def parse_args() -> argparse.Namespace:
117124 return parser .parse_args ()
118125
119126# -------------------- Dependency check --------------------
127+
128+
120129def ensure_grpc_tools () -> None :
121130 try :
122131 import grpc_tools # noqa: F401
@@ -131,6 +140,7 @@ def ensure_grpc_tools() -> None:
131140
132141# -------------------- Download & extract --------------------
133142
143+
134144def is_safe_tar_member (member : tarfile .TarInfo , base : Path ) -> bool :
135145 name = member .name
136146 if not name or name .startswith ("/" ):
@@ -145,6 +155,7 @@ def is_safe_tar_member(member: tarfile.TarInfo, base: Path) -> bool:
145155 return False
146156 return True
147157
158+
148159def safe_extract_tar_stream (response , dest : Path ) -> None :
149160 """Stream-extract a GitHub tgz, stripping the top-level folder safely."""
150161 with tarfile .open (fileobj = response , mode = "r|gz" ) as tar :
@@ -155,12 +166,15 @@ def safe_extract_tar_stream(response, dest: Path) -> None:
155166 continue
156167 if not is_safe_tar_member (member , dest ):
157168 raise RuntimeError (f"Unsafe path in archive: { member .name } " )
158- tar .extract (member , path = dest ) # nosec B202 - path validated by is_safe_tar_member
169+ # nosec B202 - path validated by is_safe_tar_member
170+ tar .extract (member , path = dest )
171+
159172
160173def validate_url_and_version (url : str , hapi_version : str ) -> None :
161174 parsed = urlparse (url )
162175 if parsed .scheme != "https" or parsed .netloc != "github.com" :
163- raise RuntimeError (f"Refusing to fetch from non-https or unexpected host: { url } " )
176+ raise RuntimeError (
177+ f"Refusing to fetch from non-https or unexpected host: { url } " )
164178
165179 if not re .fullmatch (r"v\d+\.\d+\.\d+" , hapi_version ):
166180 raise RuntimeError (f"Unexpected HAPI tag format: { hapi_version } " )
@@ -179,7 +193,7 @@ def download_and_setup_protos(hapi_version: str, protos_dir: Path) -> None:
179193 validate_url_and_version (url , hapi_version )
180194
181195 try :
182- with urllib .request .urlopen (url , timeout = 30 ) as resp : # nosec B310
196+ with urllib .request .urlopen (url , timeout = 30 ) as resp : # nosec B310
183197 safe_extract_tar_stream (resp , protos_dir )
184198 except URLError as e :
185199 raise RuntimeError (f"Failed to download protobuf files: { e } " ) from e
@@ -195,6 +209,7 @@ def download_and_setup_protos(hapi_version: str, protos_dir: Path) -> None:
195209
196210# -------------------- Filesystem helpers --------------------
197211
212+
198213def clean_and_prepare_output_dirs (* dirs : Path ) -> None :
199214 seen : set [Path ] = set ()
200215 for d in (p .resolve () for p in dirs ):
@@ -207,18 +222,21 @@ def clean_and_prepare_output_dirs(*dirs: Path) -> None:
207222 logging .info ("Creating output dir: %s" , d )
208223 d .mkdir (parents = True , exist_ok = True )
209224
225+
210226def ensure_subpackages (base : Path , subdirs : Iterable [Path ]) -> None :
211227 for p in subdirs :
212228 (base / p ).mkdir (parents = True , exist_ok = True )
213229
230+
214231def create_init_files (* roots : Path ) -> None :
215232 for root in roots :
216233 for p in [root , * root .rglob ("*" )]:
217234 if p .is_dir ():
218235 (p / "__init__.py" ).touch (exist_ok = True )
219236
237+
220238def log_generated_files (output_dir : Path ) -> None :
221- py_files = sorted (output_dir .rglob ("*.py" ))
239+ py_files = sorted (output_dir .rglob ("*.py" ))
222240 pyi_files = sorted (output_dir .rglob ("*.pyi" ))
223241
224242 print (f"\n 📂 Generated compiled proto files in { output_dir } :" )
@@ -244,6 +262,8 @@ def log_generated_files(output_dir: Path) -> None:
244262 logging .trace (" - %s" , rel )
245263
246264# -------------------- Protoc invocation --------------------
265+
266+
247267def run_protoc (
248268 proto_paths : List [Path ],
249269 out_py : Path ,
@@ -265,7 +285,7 @@ def run_protoc(
265285 args += ["--python_out" , str (out_py ), "--grpc_python_out" , str (out_grpc )]
266286 if pyi_out :
267287 args += ["--pyi_out" , str (out_py )]
268- # as_posix forces a return of the paths as / as required for protoc
288+ # as_posix forces a return of the paths as / as required for protoc
269289 args += [f .as_posix () for f in files ]
270290
271291 logging .trace ("protoc args: %s" , " " .join (args ))
@@ -276,6 +296,7 @@ def run_protoc(
276296
277297# -------------------- Import normalization for .proto --------------------
278298
299+
279300def rewrite_import_line (line : str , src_root : Path ) -> str :
280301 """
281302 Normalize an 'import "..."' line to a canonical path:
@@ -303,6 +324,7 @@ def rewrite_import_line(line: str, src_root: Path) -> str:
303324
304325 return line
305326
327+
306328def parse_import_line (line : str , src_root : Path ) -> tuple [str , Path | None ]:
307329 """Return (possibly rewritten) line and a dependency Path (or None)."""
308330 if "import " not in line or ".proto" not in line :
@@ -435,31 +457,49 @@ def compile_mirror(protos_root: Path, mirror_out: Path) -> None:
435457 shutil .rmtree (temp_root , ignore_errors = True )
436458
437459# -------------------- Post-generation Python import fixups --------------------
460+
461+
438462def _iter_py_like (root : Path ):
439463 for pat in ("*.py" , "*.pyi" ):
440464 yield from root .rglob (pat )
441465
466+
442467# -------------------- Import-rewrite helpers (precompiled regexes) --------------------
443468# Keep these at module scope so they don't count against function complexity and are compiled once.
444- _RX_IMPORT_AS = re .compile (r"^\s*import (\w+_pb2) as" , re .MULTILINE )
445- _RX_FROM_SERVICES_AS = re .compile (r"^\s*from\s+services\s+import\s+(\w+_pb2)\s+as" , re .MULTILINE )
446- _RX_FROM_SERVICES = re .compile (r"^\s*from\s+services\s+import\s+(\w+_pb2)\b" , re .MULTILINE )
447- _RX_FROM_SERVICES_SUBPKG = re .compile (r"^\s*from\s+services\.((?:\w+\.)*\w+)\s+import\s+(\w+_pb2)(\s+as\b)?" , re .MULTILINE )
448- _RX_IMPORT_SERVICES_AS = re .compile (r"^\s*import\s+services\.((?:\w+\.)*)(\w+_pb2)\s+as" , re .MULTILINE )
449- _RX_IMPORT_SERVICES = re .compile (r"^\s*import\s+services\.((?:\w+\.)*)(\w+_pb2)\b" , re .MULTILINE )
450- _RX_FROM_AUX_TSS = re .compile (r"^\s*from\s+auxiliary\.tss" , re .MULTILINE )
451- _RX_FROM_AUX_HINTS = re .compile (r"^\s*from\s+auxiliary\.hints" , re .MULTILINE )
452- _RX_FROM_AUX_HISTORY = re .compile (r"^\s*from\s+auxiliary\.history" , re .MULTILINE )
453- _RX_FROM_EVENT = re .compile (r"^\s*from\s+event\s+import\s+(\w+_pb2)(\s+as\b)?" , re .MULTILINE )
454- _RX_FROM_PLATFORM_EVENT = re .compile (r"^\s*from\s+platform\.event\s+import\s+(\w+_pb2)(\s+as\b)?" , re .MULTILINE )
455- _RX_FROM_DOT_STATE = re .compile (r"^\s*from\s+\.\s*state(\.[\w\.]+)?\s+import\s+(\w+_pb2)(\s+as\b)?" , re .MULTILINE )
456- _RX_FROM_ABS_STATE = re .compile (r"^\s*from\s+services\.state(\.[\w\.]+)?\s+import\s+(\w+_pb2)(\s+as\b)?" , re .MULTILINE )
457- _RX_FROM_DOT_IMPORT_LOCAL = re .compile (r"^\s*from\s+\.\s+import\s+(\w+_pb2)(\s+as\s+\w+)?\s*$" , re .MULTILINE )
458-
459- _RX_MIRROR_IMPORT_AS = re .compile (r"^\s*import (\w+_pb2) as" , re .MULTILINE )
460- _RX_FROM_MIRROR_AS = re .compile (r"^\s*from\s+mirror\s+import\s+(\w+_pb2)\s+as" , re .MULTILINE )
461- _RX_FROM_SERVICES_AS_MIR = re .compile (r"^\s*from\s+services\s+import\s+(\w+_pb2)\s+as" , re .MULTILINE )
462- _RX_FROM_DOT_AS_MIR = re .compile (r"^\s*from\s+\.\s+import\s+(\w+_pb2)\s+as" , re .MULTILINE )
469+ _RX_IMPORT_AS = re .compile (r"^\s*import (\w+_pb2) as" , re .MULTILINE )
470+ _RX_FROM_SERVICES_AS = re .compile (
471+ r"^\s*from\s+services\s+import\s+(\w+_pb2)\s+as" , re .MULTILINE )
472+ _RX_FROM_SERVICES = re .compile (
473+ r"^\s*from\s+services\s+import\s+(\w+_pb2)\b" , re .MULTILINE )
474+ _RX_FROM_SERVICES_SUBPKG = re .compile (
475+ r"^\s*from\s+services\.((?:\w+\.)*\w+)\s+import\s+(\w+_pb2)(\s+as\b)?" , re .MULTILINE )
476+ _RX_IMPORT_SERVICES_AS = re .compile (
477+ r"^\s*import\s+services\.((?:\w+\.)*)(\w+_pb2)\s+as" , re .MULTILINE )
478+ _RX_IMPORT_SERVICES = re .compile (
479+ r"^\s*import\s+services\.((?:\w+\.)*)(\w+_pb2)\b" , re .MULTILINE )
480+ _RX_FROM_AUX_TSS = re .compile (r"^\s*from\s+auxiliary\.tss" , re .MULTILINE )
481+ _RX_FROM_AUX_HINTS = re .compile (r"^\s*from\s+auxiliary\.hints" , re .MULTILINE )
482+ _RX_FROM_AUX_HISTORY = re .compile (
483+ r"^\s*from\s+auxiliary\.history" , re .MULTILINE )
484+ _RX_FROM_EVENT = re .compile (
485+ r"^\s*from\s+event\s+import\s+(\w+_pb2)(\s+as\b)?" , re .MULTILINE )
486+ _RX_FROM_PLATFORM_EVENT = re .compile (
487+ r"^\s*from\s+platform\.event\s+import\s+(\w+_pb2)(\s+as\b)?" , re .MULTILINE )
488+ _RX_FROM_DOT_STATE = re .compile (
489+ r"^\s*from\s+\.\s*state(\.[\w\.]+)?\s+import\s+(\w+_pb2)(\s+as\b)?" , re .MULTILINE )
490+ _RX_FROM_ABS_STATE = re .compile (
491+ r"^\s*from\s+services\.state(\.[\w\.]+)?\s+import\s+(\w+_pb2)(\s+as\b)?" , re .MULTILINE )
492+ _RX_FROM_DOT_IMPORT_LOCAL = re .compile (
493+ r"^\s*from\s+\.\s+import\s+(\w+_pb2)(\s+as\s+\w+)?\s*$" , re .MULTILINE )
494+
495+ _RX_MIRROR_IMPORT_AS = re .compile (r"^\s*import (\w+_pb2) as" , re .MULTILINE )
496+ _RX_FROM_MIRROR_AS = re .compile (
497+ r"^\s*from\s+mirror\s+import\s+(\w+_pb2)\s+as" , re .MULTILINE )
498+ _RX_FROM_SERVICES_AS_MIR = re .compile (
499+ r"^\s*from\s+services\s+import\s+(\w+_pb2)\s+as" , re .MULTILINE )
500+ _RX_FROM_DOT_AS_MIR = re .compile (
501+ r"^\s*from\s+\.\s+import\s+(\w+_pb2)\s+as" , re .MULTILINE )
502+
463503
464504def _walk_and_rewrite (root : Path , rewriter ) -> tuple [int , int ]:
465505 """Walk .py and .pyi under root, rewrite with `rewriter(text, path) -> new_text|None`."""
@@ -501,7 +541,8 @@ def rewriter(text: str, path: Path) -> str | None:
501541 s = _RX_FROM_AUX_HISTORY .sub (r"from .auxiliary.history" , s )
502542
503543 s = _RX_FROM_EVENT .sub (r"from ..platform.event import \1\2" , s )
504- s = _RX_FROM_PLATFORM_EVENT .sub (r"from ..platform.event import \1\2" , s )
544+ s = _RX_FROM_PLATFORM_EVENT .sub (
545+ r"from ..platform.event import \1\2" , s )
505546
506547 # State imports need dynamic dots
507548 def repl_dot_state (m : re .Match ) -> str :
@@ -556,6 +597,7 @@ def repl_from_dot(m: re.Match) -> str:
556597 return None if s == text else s
557598 return rewriter
558599
600+
559601def _rewrite_platform_event (text : str , _path : Path ) -> str | None :
560602 s = text
561603 s2 = s
@@ -573,16 +615,18 @@ def _rewrite_platform_event(text: str, _path: Path) -> str | None:
573615 r'from . import \1\2' , s2 )
574616 return None if s2 == s else s2
575617
618+
576619def adjust_python_imports (services_dir : Path , mirror_dir : Path ) -> None :
577620 logging .info ("Adjusting imports in services under %s" , services_dir )
578621 service_root_modules = {f .stem for f in services_dir .glob ("*_pb2.py" )}
579622 svc_changed , svc_total = _walk_and_rewrite (
580- services_dir , _rewrite_services_factory (services_dir , service_root_modules )
623+ services_dir , _rewrite_services_factory (
624+ services_dir , service_root_modules )
581625 )
582626 logging .info ("Services: rewrote %d/%d files" , svc_changed , svc_total )
583627
584628 logging .info ("Adjusting imports in mirror under %s" , mirror_dir )
585- mirror_modules = {f .stem for f in mirror_dir .rglob ("*_pb2.py" )}
629+ mirror_modules = {f .stem for f in mirror_dir .rglob ("*_pb2.py" )}
586630 service_modules = {f .stem for f in services_dir .rglob ("*_pb2.py" )}
587631 mir_changed , mir_total = _walk_and_rewrite (
588632 mirror_dir , _rewrite_mirror_factory (mirror_modules , service_modules )
@@ -593,9 +637,13 @@ def adjust_python_imports(services_dir: Path, mirror_dir: Path) -> None:
593637 pe_dir = services_dir .parent / "platform" / "event"
594638 if pe_dir .exists ():
595639 logging .info ("Adjusting imports in platform/event under %s" , pe_dir )
596- pe_changed , pe_total = _walk_and_rewrite (pe_dir , _rewrite_platform_event )
597- logging .info ("Platform/event: rewrote %d/%d files" , pe_changed , pe_total )
640+ pe_changed , pe_total = _walk_and_rewrite (
641+ pe_dir , _rewrite_platform_event )
642+ logging .info ("Platform/event: rewrote %d/%d files" ,
643+ pe_changed , pe_total )
598644# -------------------- Main --------------------
645+
646+
599647def main () -> None :
600648 args = parse_args ()
601649 setup_logging (args .verbose )
@@ -631,13 +679,14 @@ def main() -> None:
631679 ensure_subpackages (cfg .mirror_out , [Path ("mirror" )])
632680
633681 # Compile groups
634- compile_services_and_platform (cfg .protos_dir , cfg .services_out , cfg .pyi_out )
682+ compile_services_and_platform (
683+ cfg .protos_dir , cfg .services_out , cfg .pyi_out )
635684 compile_mirror (cfg .protos_dir , cfg .mirror_out )
636685 log_generated_files (cfg .mirror_out )
637686
638687 # Fix imports and make packages importable
639688 adjust_python_imports (cfg .services_out / Path ("services" ),
640- cfg .mirror_out / Path ("mirror" ))
689+ cfg .mirror_out / Path ("mirror" ))
641690 create_init_files (cfg .services_out , cfg .mirror_out )
642691
643692 print ("✅ All protobuf files have been generated and adjusted successfully!" )
0 commit comments