refactor: return a sequence of replacements

protoletariat/fdsetgen.py

@@ -11,7 +11,7 @@
 import astor
 from google.protobuf.descriptor_pb2 import FileDescriptorSet
 
-from .rewrite import ImportRewriter, build_import_rewrite
+from .rewrite import ImportRewriter, build_rewrites
 
 _PROTO_SUFFIX_PATTERN = re.compile(r"^(.+)\.proto$")
 
@@ -46,12 +46,14 @@ def fix_imports(
             # module, but they import it so we register a rewrite for the
             # current proto as a dependency of itself to handle the case
             # of services
-            rewriter.register_import_rewrite(build_import_rewrite(fd_name, fd_name))
+            for repl in build_rewrites(fd_name, fd_name):
+                rewriter.register_rewrite(repl)
+
+            # register _proto_ import rewrites
             for dep in fd.dependency:
                 dep_name = _remove_proto_suffix(dep)
-                rewriter.register_import_rewrite(
-                    build_import_rewrite(fd_name, dep_name)
-                )
+                for repl in build_rewrites(fd_name, dep_name):
+                    rewriter.register_rewrite(repl)
 
         for fd in fdset.file:
             fd_name = _remove_proto_suffix(fd.name)

protoletariat/rewrite.py

@@ -85,7 +85,7 @@ def __call__(self, node: AST) -> AST:
         return node
 
 
-def build_import_rewrite(proto: str, dep: str) -> Replacement:
+def build_rewrites(proto: str, dep: str) -> Sequence[Replacement]:
     """Construct a replacement import for `dep`.
 
     Parameters
@@ -119,7 +119,7 @@ def build_import_rewrite(proto: str, dep: str) -> Replacement:
         old = f"from {from_} import {part}_pb2 as {as_}"
         new = f"from {leading_dots}{from_} import {part}_pb2 as {as_}"
 
-    return Replacement(old=old, new=new)
+    return [Replacement(old=old, new=new)]
 
 
 class ImportRewriter(ast.NodeTransformer):
@@ -128,7 +128,7 @@ class ImportRewriter(ast.NodeTransformer):
     def __init__(self) -> None:
         self.rewrite = Rewriter()
 
-    def register_import_rewrite(self, replacement: Replacement) -> None:
+    def register_rewrite(self, replacement: Replacement) -> None:
         """Register a rewrite rule for turning `old` into `new`."""
         (old_import,) = ast.parse(replacement.old).body
         (new_import,) = ast.parse(replacement.new).body