From 22c856d64be103567a58bec4bb311bf37d8b1b2d Mon Sep 17 00:00:00 2001 From: Adrien Vannson Date: Thu, 16 Jan 2025 11:35:01 +0100 Subject: [PATCH 1/2] Move global message pool to generated code --- src/betterproto2_compiler/known_types/any.py | 6 ++++-- src/betterproto2_compiler/plugin/parser.py | 6 ++++++ src/betterproto2_compiler/templates/header.py.j2 | 7 +++++++ src/betterproto2_compiler/templates/template.py.j2 | 2 +- 4 files changed, 18 insertions(+), 3 deletions(-) diff --git a/src/betterproto2_compiler/known_types/any.py b/src/betterproto2_compiler/known_types/any.py index a6a6fd96..10580252 100644 --- a/src/betterproto2_compiler/known_types/any.py +++ b/src/betterproto2_compiler/known_types/any.py @@ -1,6 +1,8 @@ import betterproto2 from betterproto2.lib.std.google.protobuf import Any as VanillaAny +default_message_pool = betterproto2.MessagePool() # Only for typing purpose + class Any(VanillaAny): def pack(self, message: betterproto2.Message, message_pool: "betterproto2.MessagePool | None" = None) -> None: @@ -10,7 +12,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 +24,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] diff --git a/src/betterproto2_compiler/plugin/parser.py b/src/betterproto2_compiler/plugin/parser.py index fabca1bd..f2456f74 100644 --- a/src/betterproto2_compiler/plugin/parser.py +++ b/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) diff --git a/src/betterproto2_compiler/templates/header.py.j2 b/src/betterproto2_compiler/templates/header.py.j2 index b1ffdd83..c7658707 100644 --- a/src/betterproto2_compiler/templates/header.py.j2 +++ b/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 diff --git a/src/betterproto2_compiler/templates/template.py.j2 b/src/betterproto2_compiler/templates/template.py.j2 index 9a9fe92f..d5632a72 100644 --- a/src/betterproto2_compiler/templates/template.py.j2 +++ b/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 %} From c4bf7b39b0c4e5a5b3aff35eaaf56d868a119ec5 Mon Sep 17 00:00:00 2001 From: Adrien Vannson Date: Thu, 16 Jan 2025 11:37:59 +0100 Subject: [PATCH 2/2] Fix error --- src/betterproto2_compiler/known_types/any.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/betterproto2_compiler/known_types/any.py b/src/betterproto2_compiler/known_types/any.py index 10580252..7d0b8e9e 100644 --- a/src/betterproto2_compiler/known_types/any.py +++ b/src/betterproto2_compiler/known_types/any.py @@ -1,7 +1,9 @@ import betterproto2 from betterproto2.lib.std.google.protobuf import Any as VanillaAny -default_message_pool = betterproto2.MessagePool() # Only for typing purpose +# TODO put back +# default_message_pool = betterproto2.MessagePool() # Only for typing purpose +default_message_pool = ... class Any(VanillaAny):