Message wrapping

src/betterproto2_compiler/compile/importing.py

@@ -5,7 +5,7 @@
     TYPE_CHECKING,
 )
 
-from betterproto2_compiler.lib.google import protobuf as google_protobuf
+from betterproto2_compiler.known_types import WRAPPED_TYPES
 from betterproto2_compiler.settings import Settings
 
 from ..casing import safe_snake_case
@@ -14,18 +14,6 @@
 if TYPE_CHECKING:
     from ..plugin.models import PluginRequestCompiler
 
-WRAPPER_TYPES: dict[str, type] = {
-    ".google.protobuf.DoubleValue": google_protobuf.DoubleValue,
-    ".google.protobuf.FloatValue": google_protobuf.FloatValue,
-    ".google.protobuf.Int32Value": google_protobuf.Int32Value,
-    ".google.protobuf.Int64Value": google_protobuf.Int64Value,
-    ".google.protobuf.UInt32Value": google_protobuf.UInt32Value,
-    ".google.protobuf.UInt64Value": google_protobuf.UInt64Value,
-    ".google.protobuf.BoolValue": google_protobuf.BoolValue,
-    ".google.protobuf.StringValue": google_protobuf.StringValue,
-    ".google.protobuf.BytesValue": google_protobuf.BytesValue,
-}
-
 
 def parse_source_type_name(field_type_name: str, request: PluginRequestCompiler) -> tuple[str, str]:
     """
@@ -73,26 +61,18 @@ def get_type_reference(
     imports: set,
     source_type: str,
     request: PluginRequestCompiler,
-    unwrap: bool = True,
+    wrap: bool = True,
     settings: Settings,
 ) -> str:
     """
     Return a Python type name for a proto type reference. Adds the import if
     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"
-
-        elif source_type == ".google.protobuf.Timestamp":
-            return "datetime.datetime"
-
     source_package, source_type = parse_source_type_name(source_type, request)
 
+    if wrap and (source_package, source_type) in WRAPPED_TYPES:
+        return WRAPPED_TYPES[(source_package, source_type)]
+
     current_package: list[str] = package.split(".") if package else []
     py_package: list[str] = source_package.split(".") if source_package else []
     py_type: str = pythonize_class_name(source_type)

src/betterproto2_compiler/known_types/__init__.py

@@ -2,13 +2,102 @@
 
 from .any import Any
 from .duration import Duration
+from .google_values import (
+    BoolValue,
+    BytesValue,
+    DoubleValue,
+    FloatValue,
+    Int32Value,
+    Int64Value,
+    StringValue,
+    UInt32Value,
+    UInt64Value,
+)
 from .timestamp import Timestamp
 
 # For each (package, message name), lists the methods that should be added to the message definition.
 # The source code of the method is read from the `known_types` folder. If imports are needed, they can be directly added
 # to the template file: they will automatically be removed if not necessary.
 KNOWN_METHODS: dict[tuple[str, str], list[Callable]] = {
     ("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", "Timestamp"): [
+        Timestamp.from_datetime,
+        Timestamp.to_datetime,
+        Timestamp.timestamp_to_json,
+        Timestamp.from_dict,
+        Timestamp.to_dict,
+        Timestamp.from_wrapped,
+        Timestamp.to_wrapped,
+    ],
+    ("google.protobuf", "Duration"): [
+        Duration.from_timedelta,
+        Duration.to_timedelta,
+        Duration.delta_to_json,
+        Duration.from_dict,
+        Duration.to_dict,
+        Duration.from_wrapped,
+        Duration.to_wrapped,
+    ],
+    ("google.protobuf", "BoolValue"): [
+        BoolValue.from_dict,
+        BoolValue.to_dict,
+        BoolValue.from_wrapped,
+        BoolValue.to_wrapped,
+    ],
+    ("google.protobuf", "Int32Value"): [
+        Int32Value.from_dict,
+        Int32Value.to_dict,
+        Int32Value.from_wrapped,
+        Int32Value.to_wrapped,
+    ],
+    ("google.protobuf", "Int64Value"): [
+        Int64Value.from_dict,
+        Int64Value.to_dict,
+        Int64Value.from_wrapped,
+        Int64Value.to_wrapped,
+    ],
+    ("google.protobuf", "UInt32Value"): [
+        UInt32Value.from_dict,
+        UInt32Value.to_dict,
+        UInt32Value.from_wrapped,
+        UInt32Value.to_wrapped,
+    ],
+    ("google.protobuf", "UInt64Value"): [
+        UInt64Value.from_dict,
+        UInt64Value.to_dict,
+        UInt64Value.from_wrapped,
+        UInt64Value.to_wrapped,
+    ],
+    ("google.protobuf", "FloatValue"): [
+        FloatValue.from_dict,
+        FloatValue.to_dict,
+        FloatValue.from_wrapped,
+        FloatValue.to_wrapped,
+    ],
+    ("google.protobuf", "DoubleValue"): [
+        DoubleValue.from_dict,
+        DoubleValue.to_dict,
+        DoubleValue.from_wrapped,
+        DoubleValue.to_wrapped,
+    ],
+    ("google.protobuf", "StringValue"): [
+        StringValue.from_dict,
+        StringValue.to_dict,
+        StringValue.from_wrapped,
+        StringValue.to_wrapped,
+    ],
+    ("google.protobuf", "BytesValue"): [
+        BytesValue.from_dict,
+        BytesValue.to_dict,
+        BytesValue.from_wrapped,
+        BytesValue.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",
+    ("google.protobuf", "Timestamp"): "datetime.datetime",
+    ("google.protobuf", "Duration"): "datetime.timedelta",
 }

src/betterproto2_compiler/known_types/duration.py

@@ -1,4 +1,8 @@
 import datetime
+import re
+import typing
+
+import betterproto2
 
 from betterproto2_compiler.lib.google.protobuf import Duration as VanillaDuration
 
@@ -23,3 +27,44 @@ def delta_to_json(delta: datetime.timedelta) -> str:
             while len(parts[1]) not in (3, 6, 9):
                 parts[1] = f"{parts[1]}0"
         return f"{'.'.join(parts)}s"
+
+    # TODO typing
+    @classmethod
+    def from_dict(cls, value):
+        if isinstance(value, str):
+            if not re.match(r"^\d+(\.\d+)?s$", value):
+                raise ValueError(f"Invalid duration string: {value}")
+
+            seconds = float(value[:-1])
+            return Duration(seconds=int(seconds), nanos=int((seconds - int(seconds)) * 1e9))
+
+        return super().from_dict(value)
+
+    # TODO typing
+    def to_dict(
+        self,
+        *,
+        output_format: betterproto2.OutputFormat = betterproto2.OutputFormat.PROTO_JSON,
+        casing: betterproto2.Casing = betterproto2.Casing.CAMEL,
+        include_default_values: bool = False,
+    ) -> dict[str, typing.Any] | typing.Any:
+        # If the output format is PYTHON, we should have kept the wrapped type without building the real class
+        assert output_format == betterproto2.OutputFormat.PROTO_JSON
+
+        assert 0 <= self.nanos < 1e9
+
+        if self.nanos == 0:
+            return f"{self.seconds}s"
+
+        nanos = f"{self.nanos:09d}".rstrip("0")
+        if len(nanos) < 3:
+            nanos += "0" * (3 - len(nanos))
+
+        return f"{self.seconds}.{nanos}s"
+
+    @staticmethod
+    def from_wrapped(wrapped: datetime.timedelta) -> "Duration":
+        return Duration.from_timedelta(wrapped)
+
+    def to_wrapped(self) -> datetime.timedelta:
+        return self.to_timedelta()