feat: introduce module suffixes to allow adjustment of grpc generated…

… code

protoletariat/__main__.py

@@ -50,19 +50,30 @@ def _make_callback(overwrite: bool) -> Callable[[Path, str], None]:
     help="Recursively create __init__.py files under `--python-out`",
     show_default=True,
 )
+@click.option(
+    "-s",
+    "--module-suffixes",
+    type=str,
+    multiple=True,
+    default=["_pb2.py", "_pb2_grpc.py"],
+    help="Suffixes of Python modules to process",
+    show_default=True,
+)
 @click.pass_context
 def main(
     ctx: click.Context,
     python_out: Path,
     in_place: bool,
     create_package: bool,
+    module_suffixes: list[str],
 ) -> None:
     ctx.ensure_object(dict)
     ctx.obj.update(
         dict(
             python_out=python_out,
             create_package=create_package,
             overwrite_callback=_make_callback(in_place),
+            module_suffixes=module_suffixes,
         )
     )
 

protoletariat/fdsetgen.py

@@ -6,7 +6,7 @@
 import subprocess
 import tempfile
 from pathlib import Path
-from typing import Callable, Iterable
+from typing import Callable, Iterable, Sequence
 
 import astor
 from google.protobuf.descriptor_pb2 import FileDescriptorSet
@@ -29,6 +29,7 @@ def fix_imports(
         python_out: Path,
         create_package: bool,
         overwrite_callback: Callable[[Path, str], None],
+        module_suffixes: Sequence[str],
     ) -> None:
         fdset = FileDescriptorSet.FromString(self.generate_file_descriptor_set_bytes())
 
@@ -43,13 +44,15 @@ def fix_imports(
 
         # only rewrite things with dependencies
         for fd_name in (fd.name for fd in fdset.file if fd.dependency):
-            name = _PROTO_SUFFIX_PATTERN.sub(r"\1_pb2.py", fd_name)
-            python_file = python_out.joinpath(name)
-            raw_code = python_file.read_text()
-            module = ast.parse(raw_code)
-            new_module = rewriter.visit(module)
-            new_code = astor.to_source(new_module)
-            overwrite_callback(python_file, new_code)
+            for suffix in module_suffixes:
+                name = _PROTO_SUFFIX_PATTERN.sub(rf"\1{suffix}", fd_name)
+                python_file = python_out.joinpath(name)
+                if python_file.exists():
+                    raw_code = python_file.read_text()
+                    module = ast.parse(raw_code)
+                    new_module = rewriter.visit(module)
+                    new_code = astor.to_source(new_module)
+                    overwrite_callback(python_file, new_code)
 
         if create_package:
             python_out.joinpath("__init__.py").touch(exist_ok=True)