Add to_dict for structs

Author: AdrienVannson

src/betterproto2_compiler/known_types/__init__.py

@@ -13,13 +13,19 @@
     UInt32Value,
     UInt64Value,
 )
+from .struct import ListValue, Struct, Value
 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],
+    # Struct
+    ("google.protobuf", "Struct"): [Struct.to_dict],
+    ("google.protobuf", "Value"): [Value.to_dict],
+    ("google.protobuf", "ListValue"): [ListValue.to_dict],
+    # Time
     ("google.protobuf", "Timestamp"): [
         Timestamp.from_datetime,
         Timestamp.to_datetime,
@@ -38,6 +44,7 @@
         Duration.from_wrapped,
         Duration.to_wrapped,
     ],
+    # Values
     ("google.protobuf", "BoolValue"): [
         BoolValue.from_dict,
         BoolValue.to_dict,
@@ -96,6 +103,10 @@
 
 # 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] = {
+    # Time
+    ("google.protobuf", "Timestamp"): "datetime.datetime",
+    ("google.protobuf", "Duration"): "datetime.timedelta",
+    # Values
     ("google.protobuf", "BoolValue"): "bool",
     ("google.protobuf", "Int32Value"): "int",
     ("google.protobuf", "Int64Value"): "int",
@@ -105,6 +116,4 @@
     ("google.protobuf", "DoubleValue"): "float",
     ("google.protobuf", "StringValue"): "str",
     ("google.protobuf", "BytesValue"): "bytes",
-    ("google.protobuf", "Timestamp"): "datetime.datetime",
-    ("google.protobuf", "Duration"): "datetime.timedelta",
 }

src/betterproto2_compiler/known_types/struct.py

@@ -0,0 +1,71 @@
+import typing
+
+import betterproto2
+
+from betterproto2_compiler.lib.google.protobuf import (
+    ListValue as VanillaListValue,
+    NullValue,
+    Struct as VanillaStruct,
+    Value as VanillaValue,
+)
+
+
+class Struct(VanillaStruct):
+    # 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 wraped type without building the real class
+        assert output_format == betterproto2.OutputFormat.PROTO_JSON
+
+        json = {}
+        for name, value in self.fields.items():
+            json[name] = value.to_dict(casing=casing)
+        return json
+
+
+class Value(VanillaValue):
+    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 wraped type without building the real class
+        assert output_format == betterproto2.OutputFormat.PROTO_JSON
+
+        match self:
+            case Value(null_value=NullValue()):
+                return None
+            case Value(number_value=float(number_value)):
+                return number_value
+            case Value(string_value=str(string_value)):
+                return string_value
+            case Value(bool_value=bool(bool_value)):
+                return bool_value
+            case Value(struct_value=struct_value) if struct_value is not None:
+                return struct_value.to_dict(casing=casing)
+            case Value(list_value=list_value) if list_value is not None:
+                return list_value.to_dict(casing=casing)
+
+
+class ListValue(VanillaListValue):
+    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 wraped type without building the real class
+        assert output_format == betterproto2.OutputFormat.PROTO_JSON
+
+        json = []
+        for value in self.values:
+            json.append(value.to_dict(casing=casing))
+        return json