From 1c7a5724e0ab859a4bfbb12af288d35cf6e41be6 Mon Sep 17 00:00:00 2001 From: Adrien Vannson Date: Mon, 5 May 2025 11:24:40 +0200 Subject: [PATCH 1/2] Add to_dict for structs --- .../known_types/__init__.py | 13 +++- .../known_types/struct.py | 71 +++++++++++++++++++ 2 files changed, 82 insertions(+), 2 deletions(-) create mode 100644 src/betterproto2_compiler/known_types/struct.py 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..6384dc4b --- /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(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 From bf413c46aac8325b3dd18a5459147122b1d9069a Mon Sep 17 00:00:00 2001 From: Adrien Vannson Date: Mon, 5 May 2025 11:35:24 +0200 Subject: [PATCH 2/2] Remove useless parameter --- src/betterproto2_compiler/known_types/struct.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/betterproto2_compiler/known_types/struct.py b/src/betterproto2_compiler/known_types/struct.py index 6384dc4b..67528bf8 100644 --- a/src/betterproto2_compiler/known_types/struct.py +++ b/src/betterproto2_compiler/known_types/struct.py @@ -24,7 +24,7 @@ def to_dict( json = {} for name, value in self.fields.items(): - json[name] = value.to_dict(casing=casing) + json[name] = value.to_dict() return json @@ -49,9 +49,9 @@ def to_dict( 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) + return struct_value.to_dict() case Value(list_value=list_value) if list_value is not None: - return list_value.to_dict(casing=casing) + return list_value.to_dict() class ListValue(VanillaListValue): @@ -67,5 +67,5 @@ def to_dict( json = [] for value in self.values: - json.append(value.to_dict(casing=casing)) + json.append(value.to_dict()) return json