def populate_fields(proto_file, root: ProtoNode) -> None:
"""Traverses a proto file, adding all message and enum fields to a tree."""
def populate_message(node, message):
"""Recursively populates nested messages and enums."""
add_message_fields(root, node, message)
for enum in message.enum_type:
add_enum_fields(node.find(enum.name), enum)
for msg in message.nested_type:
populate_message(node.find(msg.name), msg)
# Iterate through the proto file, populating top-level enums and messages.
for enum in proto_file.enum_type:
add_enum_fields(root.find(enum.name), enum)
for message in proto_file.message_type:
populate_message(root.find(message.name), message)
def add_message_fields(
root: ProtoNode,
message: ProtoNode,
proto_message,
) -> None:
"""Adds fields from a protobuf message descriptor to a message node."""
assert message.type() == ProtoNode.Type.MESSAGE
for field in proto_message.field:
if field.type_name:
# The "type_name" member contains the global .proto path of the
# field's type object, for example ".pw.protobuf.test.KeyValuePair".
# Since only a single proto file is currently supported, the root
# node has the value of the file's package ("pw.protobuf.test").
# This must be stripped from the path to find the desired node
# within the tree.
#
# TODO(frolv): Once multiple files are supported, the root node
# should refer to the global namespace, and this should no longer
# be needed.
path = field.type_name
if path[0] == '.':
path = path[1:]
if path.startswith(root.name()):
relative_path = path[len(root.name()):].lstrip('.')
else:
relative_path = path
type_node = root.find(relative_path)
else:
type_node = None
repeated = \
field.label == descriptor_pb2.FieldDescriptorProto.LABEL_REPEATED
message.add_field(
ProtoMessageField(
field.name,
field.number,
field.type,
type_node,
repeated,
))