Add missing file

Author: AdrienVannson

betterproto2_compiler/src/betterproto2_compiler/known_types/struct.py

@@ -0,0 +1,131 @@
+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
+    @classmethod
+    def from_dict(cls, value, *, ignore_unknown_fields: bool = False):
+        assert isinstance(value, dict)
+
+        fields: dict[str, Value] = {}
+
+        for key, val in value.items():
+            fields[key] = Value.from_dict(val, ignore_unknown_fields=ignore_unknown_fields)
+
+        return cls(fields=fields)
+
+    # 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
+
+        return {
+            key: value.to_dict(
+                output_format=output_format, casing=casing, include_default_values=include_default_values
+            )
+            for key, value in self.fields.items()
+        }
+
+    @staticmethod
+    def from_wrapped(wrapped: betterproto2.JSON) -> "Struct":
+        return Struct.from_dict(wrapped)
+
+    def to_wrapped(self) -> betterproto2.JSON:
+        return self.to_dict()
+
+
+class Value(VanillaValue):
+    # TODO typing
+    @classmethod
+    def from_dict(cls, value, *, ignore_unknown_fields: bool = False):
+        match value:
+            case bool() as b:
+                return cls(bool_value=b)
+            case int() | float() as num:
+                return cls(number_value=num)
+            case str() as s:
+                return cls(string_value=s)
+            case list() as l:
+                return cls(list_value=ListValue.from_dict(l))
+            case dict() as d:
+                return cls(struct_value=Struct.from_dict(d))
+            case None:
+                return cls(null_value=NullValue.NULL_VALUE)
+            case _:
+                raise ValueError(f"Unknown value type: {type(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
+
+        match self:
+            case Value(null_value=NullValue.NULL_VALUE):
+                return None
+            case Value(bool_value=bool(b)):
+                return b
+            case Value(number_value=int(num)) | Value(number_value=float(num)):
+                return num
+            case Value(string_value=str(s)):
+                return s
+            case Value(list_value=ListValue(values=l)):
+                return [v.to_dict() for v in l]
+            case Value(struct_value=Struct(fields=f)):
+                return {k: v.to_dict() for k, v in f.items()}
+
+        raise ValueError("Invalid value")
+
+    @staticmethod
+    def from_wrapped(wrapped: betterproto2.JSON) -> "Value":
+        return Value.from_dict(wrapped)
+
+    def to_wrapped(self) -> betterproto2.JSON:
+        return self.to_dict()
+
+
+class ListValue(VanillaListValue):
+    # TODO typing
+    @classmethod
+    def from_dict(cls, value, *, ignore_unknown_fields: bool = False):
+        return cls(values=[Value.from_dict(v) for v in 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
+
+        return [value.to_dict() for value in self.values]
+
+    @staticmethod
+    def from_wrapped(wrapped: list[betterproto2.JSON]) -> "ListValue":
+        return ListValue.from_dict(wrapped)
+
+    def to_wrapped(self) -> list[betterproto2.JSON]:
+        return self.to_dict()