Change wrapping mechanism

Author: AdrienVannson

src/betterproto2_compiler/compile/importing.py

@@ -81,10 +81,6 @@ def get_type_reference(
     necessary. Unwraps well known type if required.
     """
     if unwrap:
-        if source_type in WRAPPER_TYPES:
-            wrapped_type = type(WRAPPER_TYPES[source_type]().value)
-            return f"{wrapped_type.__name__} | None"
-
         if source_type == ".google.protobuf.Duration":
             return "datetime.timedelta"
 

src/betterproto2_compiler/known_types/__init__.py

@@ -2,6 +2,7 @@
 
 from .any import Any
 from .duration import Duration
+from .google_values import BoolValue, StringValue
 from .timestamp import Timestamp
 
 # For each (package, message name), lists the methods that should be added to the message definition.
@@ -11,4 +12,12 @@
     ("google.protobuf", "Any"): [Any.pack, Any.unpack, Any.to_dict],
     ("google.protobuf", "Timestamp"): [Timestamp.from_datetime, Timestamp.to_datetime, Timestamp.timestamp_to_json],
     ("google.protobuf", "Duration"): [Duration.from_timedelta, Duration.to_timedelta, Duration.delta_to_json],
+    ("google.protobuf", "BoolValue"): [BoolValue.from_wrapped, BoolValue.to_wrapped],
+    ("google.protobuf", "StringValue"): [StringValue.from_wrapped, StringValue.to_wrapped],
+}
+
+# A wrapped type is the type of a message that is automatically replaced by a known Python type.
+WRAPPED_TYPES: dict[tuple[str, str], str] = {
+    ("google.protobuf", "BoolValue"): "bool",
+    ("google.protobuf", "StringValue"): "str",
 }

src/betterproto2_compiler/plugin/models.py

@@ -26,14 +26,12 @@
 
 import builtins
 import inspect
-import re
 from collections.abc import Iterator
 from dataclasses import (
     dataclass,
     field,
 )
 
-import betterproto2
 from betterproto2 import unwrap
 
 from betterproto2_compiler.compile.importing import get_type_reference, parse_source_type_name
@@ -43,7 +41,7 @@
     pythonize_field_name,
     pythonize_method_name,
 )
-from betterproto2_compiler.known_types import KNOWN_METHODS
+from betterproto2_compiler.known_types import KNOWN_METHODS, WRAPPED_TYPES
 from betterproto2_compiler.lib.google.protobuf import (
     DescriptorProto,
     EnumDescriptorProto,
@@ -318,8 +316,15 @@ def get_field_string(self) -> str:
     @property
     def betterproto_field_args(self) -> list[str]:
         args = []
-        if self.field_wraps:
-            args.append(f"wraps={self.field_wraps}")
+
+        if self.proto_obj.type == FieldDescriptorProtoType.TYPE_MESSAGE:
+            type_package, type_name = parse_source_type_name(self.proto_obj.type_name, self.output_file.parent_request)
+
+            if (type_package, type_name) in WRAPPED_TYPES:
+                # Without the lambda function, the type is evaluated right away, which fails since the corresponding
+                # import is placed at the end of the file to avoid circular imports.
+                args.append(f"wrap=lambda: {self.py_type}")
+
         if self.optional:
             args.append("optional=True")
         elif self.repeated:
@@ -338,16 +343,6 @@ def use_builtins(self) -> bool:
             self.py_type == self.py_name and self.py_name in dir(builtins)
         )
 
-    @property
-    def field_wraps(self) -> str | None:
-        """Returns betterproto wrapped field type or None."""
-        match_wrapper = re.match(r"\.google\.protobuf\.(.+)Value$", self.proto_obj.type_name)
-        if match_wrapper:
-            wrapped_type = "TYPE_" + match_wrapper.group(1).upper()
-            if hasattr(betterproto2, wrapped_type):
-                return f"betterproto2.{wrapped_type}"
-        return None
-
     @property
     def repeated(self) -> bool:
         return self.proto_obj.label == FieldDescriptorProtoLabel.LABEL_REPEATED
@@ -405,6 +400,14 @@ def py_type(self) -> str:
     @property
     def annotation(self) -> str:
         py_type = self.py_type
+
+        # Replace by the wrapping type if needed
+        if self.proto_obj.type == FieldDescriptorProtoType.TYPE_MESSAGE:
+            type_package, type_name = parse_source_type_name(self.proto_obj.type_name, self.output_file.parent_request)
+
+            if wrapped_type := WRAPPED_TYPES.get((type_package, type_name)):
+                py_type = wrapped_type
+
         if self.use_builtins:
             py_type = f"builtins.{py_type}"
         if self.repeated: