protoc-gen-python_grpc_typings

#!/usr/bin/python3
import sys

from google.protobuf.compiler import plugin_pb2
from google.protobuf.descriptor_pb2 import FileDescriptorProto

from stubs_generator.base import ConstantPart, NEW_LINE
from stubs_generator.messages import File, Import
from stubs_generator.servicers import AbstractMethod, AddToServerMethod, Servicer, Stub, StubMethod
from stubs_generator.utils import ImportPool, before_every, decode_type, get_comments

DEFAULT_TAB_STR = '    '


def generate_pb2_grpc_stub_file_content(proto_descriptor: FileDescriptorProto) -> str:
    """Generates typing stub file for messages"""
    comments = get_comments(proto_file)
    import_pool = ImportPool()
    import_pool.add(Import('grpc', ['ServicerContext', 'Channel', 'Server', 'CallCredentials']))
    import_pool.add(Import('abc', ['ABC', 'abstractmethod']))
    import_pool.add(Import('typing', ['Any']))
    return File(
        # Header for a file
        ConstantPart("""\
# ############################################################################# #
#  Automatically generated protobuf stub files for python                       #
#   by protoc-gen-python_grpc_typings plugin for protoc                         #
# ############################################################################# #

"""),
        import_pool,
        # Stub servicer
        *before_every(
            [NEW_LINE, NEW_LINE],
            *[Stub(
                s.name,
                *[StubMethod(meth.name,
                             decode_type(name=meth.input_type,
                                         import_pool=import_pool,
                                         proto_name=proto_descriptor.name[:-6]),
                             decode_type(name=meth.output_type,
                                         import_pool=import_pool,
                                         proto_name=proto_descriptor.name[:-6]),
                             comments.get("{}.{}".format(s.name, meth.name), []))
                  for meth in s.method],
                comments=comments.get(s.name, []),
            ) for s in proto_descriptor.service],
            # Abstract servicer
            *[Servicer(
                s.name,
                *[AbstractMethod(meth.name,
                                 decode_type(name=meth.input_type,
                                             import_pool=import_pool,
                                             proto_name=proto_descriptor.name[:-6]),
                                 decode_type(name=meth.output_type,
                                             import_pool=import_pool,
                                             proto_name=proto_descriptor.name[:-6]),
                                 comments.get("{}.{}".format(s.name, meth.name), []))
                  for meth in s.method],
                comments=comments.get(s.name, []),
            ) for s in proto_descriptor.service],
            *[AddToServerMethod(s.name)
              for s in proto_descriptor.service]
        )
    ).generate(0, DEFAULT_TAB_STR)


if __name__ == '__main__':
    # Read request message from stdin
    data = sys.stdin.buffer.read()

    # Parse request
    request = plugin_pb2.CodeGeneratorRequest()
    request.ParseFromString(data)

    # Create response
    response = plugin_pb2.CodeGeneratorResponse()

    for proto_file in request.proto_file:
        if proto_file.name in request.file_to_generate:
            response.file.add(
                name="{}_pb2_grpc.pyi".format(proto_file.name[:-6]),
                content=generate_pb2_grpc_stub_file_content(proto_file)
            )
            # response.file.add(
            #     name="debug_for_{}_proto".format(proto_file.name[:-6]),
            #     content=str(proto_file)
            # )

    # Serialise response message
    output = response.SerializeToString()

    # Write to stdout
    sys.stdout.buffer.write(output)