proto_generator.py

"""
proto_generator.py - .proto 文件生成器
从 FileDescriptorProto 或 protobuf-c 结构体信息生成 .proto 文件
"""

from google.protobuf.descriptor_pb2 import (
    FileDescriptorProto, DescriptorProto, FieldDescriptorProto,
    EnumDescriptorProto, OneofDescriptorProto, ServiceDescriptorProto,
    MethodDescriptorProto,
)


# protobuf wire type name mapping
FIELD_TYPE_NAMES = {
    FieldDescriptorProto.TYPE_DOUBLE: 'double',
    FieldDescriptorProto.TYPE_FLOAT: 'float',
    FieldDescriptorProto.TYPE_INT64: 'int64',
    FieldDescriptorProto.TYPE_UINT64: 'uint64',
    FieldDescriptorProto.TYPE_INT32: 'int32',
    FieldDescriptorProto.TYPE_FIXED64: 'fixed64',
    FieldDescriptorProto.TYPE_FIXED32: 'fixed32',
    FieldDescriptorProto.TYPE_BOOL: 'bool',
    FieldDescriptorProto.TYPE_STRING: 'string',
    FieldDescriptorProto.TYPE_GROUP: 'group',
    FieldDescriptorProto.TYPE_MESSAGE: 'message',
    FieldDescriptorProto.TYPE_BYTES: 'bytes',
    FieldDescriptorProto.TYPE_UINT32: 'uint32',
    FieldDescriptorProto.TYPE_ENUM: 'enum',
    FieldDescriptorProto.TYPE_SFIXED32: 'sfixed32',
    FieldDescriptorProto.TYPE_SFIXED64: 'sfixed64',
    FieldDescriptorProto.TYPE_SINT32: 'sint32',
    FieldDescriptorProto.TYPE_SINT64: 'sint64',
}

LABEL_NAMES = {
    FieldDescriptorProto.LABEL_OPTIONAL: 'optional',
    FieldDescriptorProto.LABEL_REQUIRED: 'required',
    FieldDescriptorProto.LABEL_REPEATED: 'repeated',
}

# well-known proto files to skip
WELL_KNOWN_PROTOS = {
    'google/protobuf/any.proto',
    'google/protobuf/api.proto',
    'google/protobuf/compiler/plugin.proto',
    'google/protobuf/descriptor.proto',
    'google/protobuf/duration.proto',
    'google/protobuf/empty.proto',
    'google/protobuf/field_mask.proto',
    'google/protobuf/source_context.proto',
    'google/protobuf/struct.proto',
    'google/protobuf/timestamp.proto',
    'google/protobuf/type.proto',
    'google/protobuf/wrappers.proto',
    'google/protobuf/cpp_features.proto',
    'google/protobuf/java_features.proto',
    'google/protobuf/go_features.proto',
}


def is_well_known(filename: str) -> bool:
    """判断是否为 protobuf well-known type"""
    normalized = filename.replace('\\', '/')
    return normalized in WELL_KNOWN_PROTOS or normalized.startswith('google/protobuf/')


class ProtoFileGenerator:
    """从 FileDescriptorProto 生成 .proto 文件内容"""

    def __init__(self, fd: FileDescriptorProto, syntax: str = None):
        self.fd = fd
        # 自动推断 syntax
        if syntax:
            self.syntax = syntax
        elif fd.HasField('syntax') and fd.syntax:
            self.syntax = fd.syntax
        else:
            self.syntax = 'proto3'

    def generate(self) -> str:
        lines = []
        lines.append(f'syntax = "{self.syntax}";')
        lines.append('')

        if self.fd.package:
            lines.append(f'package {self.fd.package};')
            lines.append('')

        # imports
        for dep in self.fd.dependency:
            lines.append(f'import "{dep}";')
        if self.fd.dependency:
            lines.append('')

        # options
        opts = []
        if self.fd.HasField('options'):
            if self.fd.options.HasField('java_package'):
                opts.append(f'option java_package = "{self.fd.options.java_package}";')
            if self.fd.options.HasField('go_package'):
                opts.append(f'option go_package = "{self.fd.options.go_package}";')
        if opts:
            lines.extend(opts)
            lines.append('')

        # enums (top-level)
        for enum_desc in self.fd.enum_type:
            lines.extend(self._gen_enum(enum_desc, 0))
            lines.append('')

        # messages
        for msg_desc in self.fd.message_type:
            lines.extend(self._gen_message(msg_desc, 0))
            lines.append('')

        # services
        for svc_desc in self.fd.service:
            lines.extend(self._gen_service(svc_desc, 0))
            lines.append('')

        return '\n'.join(lines)

    def _indent(self, level: int) -> str:
        return '  ' * level

    def _gen_message(self, msg: DescriptorProto, level: int) -> list:
        ind = self._indent(level)
        lines = [f'{ind}message {msg.name} {{']

        # nested enums
        for enum_desc in msg.enum_type:
            lines.extend(self._gen_enum(enum_desc, level + 1))

        # nested messages
        for nested in msg.nested_type:
            # 跳过 map entry 类型
            if nested.options and nested.options.map_entry:
                continue
            lines.extend(self._gen_message(nested, level + 1))

        # oneofs
        oneof_fields = {}  # oneof_index -> [fields]
        for field in msg.field:
            if field.HasField('oneof_index'):
                oneof_fields.setdefault(field.oneof_index, []).append(field)

        # regular fields (non-oneof)
        for field in msg.field:
            if field.HasField('oneof_index'):
                continue
            lines.append(self._gen_field(field, level + 1))

        # oneof declarations
        for idx, oneof in enumerate(msg.oneof_decl):
            fields = oneof_fields.get(idx, [])
            ind2 = self._indent(level + 1)
            lines.append(f'{ind2}oneof {oneof.name} {{')
            for field in fields:
                lines.append(self._gen_field(field, level + 2, in_oneof=True))
            lines.append(f'{ind2}}}')

        # map fields
        for nested in msg.nested_type:
            if nested.options and nested.options.map_entry:
                key_field = None
                value_field = None
                for f in nested.field:
                    if f.number == 1:
                        key_field = f
                    elif f.number == 2:
                        value_field = f
                if key_field and value_field:
                    ind2 = self._indent(level + 1)
                    key_type = self._field_type_name(key_field)
                    val_type = self._field_type_name(value_field)
                    # 找到引用这个 map 的 field
                    for field in msg.field:
                        if field.type_name.endswith(nested.name):
                            lines.append(f'{ind2}map<{key_type}, {val_type}> {field.name} = {field.number};')

        # reserved
        for rr in msg.reserved_range:
            lines.append(f'{self._indent(level + 1)}reserved {rr.start} to {rr.end - 1};')
        for rn in msg.reserved_name:
            lines.append(f'{self._indent(level + 1)}reserved "{rn}";')

        lines.append(f'{ind}}}')
        return lines

    def _gen_field(self, field: FieldDescriptorProto, level: int, in_oneof: bool = False) -> str:
        ind = self._indent(level)
        type_name = self._field_type_name(field)

        label = ''
        if not in_oneof:
            if self.syntax == 'proto2':
                label = LABEL_NAMES.get(field.label, '') + ' '
            elif field.label == FieldDescriptorProto.LABEL_REPEATED:
                label = 'repeated '
            elif field.label == FieldDescriptorProto.LABEL_OPTIONAL and field.HasField('proto3_optional') and field.proto3_optional:
                label = 'optional '

        return f'{ind}{label}{type_name} {field.name} = {field.number};'

    def _field_type_name(self, field: FieldDescriptorProto) -> str:
        if field.type in (FieldDescriptorProto.TYPE_MESSAGE, FieldDescriptorProto.TYPE_ENUM):
            # 使用 type_name, 去掉前缀的点
            name = field.type_name
            if name.startswith('.'):
                name = name[1:]
            return name
        return FIELD_TYPE_NAMES.get(field.type, f'unknown_type_{field.type}')

    def _gen_enum(self, enum: EnumDescriptorProto, level: int) -> list:
        ind = self._indent(level)
        lines = [f'{ind}enum {enum.name} {{']
        for val in enum.value:
            lines.append(f'{self._indent(level + 1)}{val.name} = {val.number};')
        lines.append(f'{ind}}}')
        return lines

    def _gen_service(self, svc: ServiceDescriptorProto, level: int) -> list:
        ind = self._indent(level)
        lines = [f'{ind}service {svc.name} {{']
        for method in svc.method:
            input_type = method.input_type.lstrip('.')
            output_type = method.output_type.lstrip('.')
            cs = 'stream ' if method.client_streaming else ''
            ss = 'stream ' if method.server_streaming else ''
            lines.append(f'{self._indent(level + 1)}rpc {method.name} ({cs}{input_type}) returns ({ss}{output_type});')
        lines.append(f'{ind}}}')
        return lines


def fdproto_to_proto_string(fd: FileDescriptorProto) -> str:
    """快捷函数：FileDescriptorProto → .proto 文件内容"""
    gen = ProtoFileGenerator(fd)
    return gen.generate()