Move global message pool to generated code

src/betterproto2_compiler/known_types/any.py

@@ -1,6 +1,10 @@
 import betterproto2
 from betterproto2.lib.std.google.protobuf import Any as VanillaAny
 
+# TODO put back
+# default_message_pool = betterproto2.MessagePool()  # Only for typing purpose
+default_message_pool = ...
+
 
 class Any(VanillaAny):
     def pack(self, message: betterproto2.Message, message_pool: "betterproto2.MessagePool | None" = None) -> None:
@@ -10,7 +14,7 @@ def pack(self, message: betterproto2.Message, message_pool: "betterproto2.Messag
         The message type must be registered in the message pool, which is done automatically when the module defining
         the message type is imported.
         """
-        message_pool = message_pool or betterproto2.default_message_pool
+        message_pool = message_pool or default_message_pool
 
         self.type_url = message_pool.type_to_url[type(message)]
         self.value = bytes(message)
@@ -22,7 +26,7 @@ def unpack(self, message_pool: "betterproto2.MessagePool | None" = None) -> bett
         The target message type must be registered in the message pool, which is done automatically when the module
         defining the message type is imported.
         """
-        message_pool = message_pool or betterproto2.default_message_pool
+        message_pool = message_pool or default_message_pool
 
         message_type = message_pool.url_to_type[self.type_url]
 

src/betterproto2_compiler/plugin/parser.py

@@ -173,6 +173,12 @@ def generate_code(request: CodeGeneratorRequest) -> CodeGeneratorResponse:
     for init_file in init_files:
         response.file.append(CodeGeneratorResponseFile(name=str(init_file)))
 
+    response.file.append(
+        CodeGeneratorResponseFile(
+            name="message_pool.py", content="import betterproto2\n\ndefault_message_pool = betterproto2.MessagePool()\n"
+        )
+    )
+
     for output_package_name in sorted(output_paths.union(init_files)):
         print(f"Writing {output_package_name}", file=sys.stderr)
 

src/betterproto2_compiler/templates/header.py.j2

@@ -44,6 +44,13 @@ import grpclib
 
 from typing import TYPE_CHECKING
 
+{# Import the message pool of the generated code. #}
+{% if output_file.package %}
+from {{ "." * output_file.package.count(".") }}..message_pool import default_message_pool
+{% else %}
+from .message_pool import default_message_pool
+{% endif %}
+
 if TYPE_CHECKING:
     import grpclib.server
     from betterproto2.grpc.grpclib_client import MetadataLike

src/betterproto2_compiler/templates/template.py.j2

@@ -81,7 +81,7 @@ class {{ message.py_name }}(betterproto2.Message):
     {{ method_source }}
     {% endfor %}
 
-betterproto2.default_message_pool.register_message("{{ output_file.package }}", "{{ message.proto_name }}", {{ message.py_name }})
+default_message_pool.register_message("{{ output_file.package }}", "{{ message.proto_name }}", {{ message.py_name }})
 
 
 {% endfor %}