diff --git a/src/betterproto2_compiler/known_types/__init__.py b/src/betterproto2_compiler/known_types/__init__.py index f6accdd6..82a58a2f 100644 --- a/src/betterproto2_compiler/known_types/__init__.py +++ b/src/betterproto2_compiler/known_types/__init__.py @@ -13,6 +13,7 @@ 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. @@ -20,6 +21,11 @@ # 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", } diff --git a/src/betterproto2_compiler/known_types/struct.py b/src/betterproto2_compiler/known_types/struct.py new file mode 100644 index 00000000..67528bf8 --- /dev/null +++ b/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() + 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() + case Value(list_value=list_value) if list_value is not None: + return list_value.to_dict() + + +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()) + return json