Skip to content
This repository was archived by the owner on Jun 9, 2025. It is now read-only.

Add back struct support #79

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 11 additions & 2 deletions src/betterproto2_compiler/known_types/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -38,6 +44,7 @@
Duration.from_wrapped,
Duration.to_wrapped,
],
# Values
("google.protobuf", "BoolValue"): [
BoolValue.from_dict,
BoolValue.to_dict,
Expand Down Expand Up @@ -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",
Expand All @@ -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",
}
71 changes: 71 additions & 0 deletions src/betterproto2_compiler/known_types/struct.py
Original file line number Diff line number Diff line change
@@ -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
Loading