Reflect Proto Messages

reflect_proto.py

from pprint import pprint

from google.protobuf.pyext.cpp_message import GeneratedProtocolMessageType
from google.protobuf.descriptor import FieldDescriptor

from myproto import schema_pb2

pb_types = {
    FieldDescriptor.TYPE_BOOL: bool,
    FieldDescriptor.TYPE_BYTES: bytes,
    FieldDescriptor.TYPE_DOUBLE: float,
    FieldDescriptor.TYPE_FLOAT: float,
    FieldDescriptor.TYPE_INT32: int,
    FieldDescriptor.TYPE_INT64: int,
    FieldDescriptor.TYPE_STRING: str,

}

modules = [dschema_pb2]


def extract_message(msg):
    fields = []
    for field in msg.DESCRIPTOR.fields:
        field_dict = dict(
            name=field.name, number=field.number, full_name=field.full_name
        )

        is_repeated = field.label == field.LABEL_REPEATED
        is_enum = field.type == FieldDescriptor.TYPE_ENUM
        is_msg = field.type == FieldDescriptor.TYPE_MESSAGE

        field_dict["repeated"] = is_repeated
        field_dict["enum"] = is_enum

        if is_msg:
            field_dict["type_name"] = field.message_type.name
        elif is_enum:
            field_dict["type_name"] = field.enum_type.name
        else:
            field_dict["type_name"] = pb_types[field.type]

        fields.append(field_dict)

    return dict(
        name=msg.DESCRIPTOR.name,
        full_name=msg.DESCRIPTOR.full_name,
        package=msg.DESCRIPTOR.file.package,
        fields=fields,
    )


def extract_module(module):
    messages = []
    for name in dir(module):
        entity = getattr(module, name)
        if isinstance(entity, GeneratedProtocolMessageType):
            msg = extract_message(entity)
            messages.append(msg)
    return messages


for module in modules:
    messages = extract_module(module)
    pprint(messages, indent=2, width=2)

"""
# schema.proto

package "schema
message User {
  string name = 1;
}


{ 
  'full_name': 'schema.User',
  'name': 'User',
  'package': 'schema'}
  'fields': [
    {
      'enum': False,
      'full_name': 'schema.User.name',
      'name': 'name',
      'number': 1,
      'repeated': False,
      'type_name': <class 'str'>},
  ]
}
"""