main.py

import os
import re
import sys

from typing import List, Any, Collection, Iterator, NamedTuple, cast
from typing import Dict, Tuple, Optional, Deque
from contextlib import contextmanager
from collections import deque

from google.protobuf.descriptor_pb2 import FileDescriptorProto, DescriptorProto
from google.protobuf.compiler.plugin_pb2 import CodeGeneratorRequest
from google.protobuf.compiler.plugin_pb2 import CodeGeneratorResponse

from grpclib import const
from grpclib import client
from grpclib import server
from grpclib import exceptions


_CARDINALITY = {
    (False, False): const.Cardinality.UNARY_UNARY,
    (True, False): const.Cardinality.STREAM_UNARY,
    (False, True): const.Cardinality.UNARY_STREAM,
    (True, True): const.Cardinality.STREAM_STREAM,
}


class Method(NamedTuple):
    name: str
    cardinality: const.Cardinality
    request_type: str
    reply_type: str


class Service(NamedTuple):
    name: str
    methods: List[Method]


class Buffer:
    def __init__(self) -> None:
        self._lines: List[str] = []
        self._indent = 0

    def add(self, string: str, *args: Any, **kwargs: Any) -> None:
        line = " " * self._indent * 4 + string.format(*args, **kwargs)
        self._lines.append(line.rstrip(" "))

    @contextmanager
    def indent(self) -> Iterator[None]:
        self._indent += 1
        try:
            yield
        finally:
            self._indent -= 1

    def content(self) -> str:
        return "\n".join(self._lines) + "\n"


def render(
    proto_file: str,
    package: str,
    imports: Collection[str],
    services: Collection[Service],
) -> str:
    buf = Buffer()
    buf.add("# Generated by the Protocol Buffers compiler. DO NOT EDIT!")
    buf.add("# source: {}", proto_file)
    buf.add("# plugin: {}", __name__)
    if not services:
        return buf.content()

    buf.add("import abc")
    buf.add("import typing")
    buf.add("")
    buf.add("import {}", const.__name__)
    buf.add("import {}", client.__name__)
    buf.add("import {}", exceptions.__name__)
    buf.add("if typing.TYPE_CHECKING:")
    with buf.indent():
        buf.add("import {}", server.__name__)

    buf.add("")
    for mod in imports:
        buf.add("import {}", mod)
    for service in services:
        if package:
            service_name = "{}.{}".format(package, service.name)
        else:
            service_name = service.name
        buf.add("")
        buf.add("")
        buf.add("class {}Base(abc.ABC):", service.name)
        with buf.indent():
            for name, _, request_type, reply_type in service.methods:
                buf.add("")
                buf.add("@abc.abstractmethod")
                buf.add(
                    "async def {}(self, stream: '{}.{}[{}, {}]') -> None:",
                    name,
                    server.__name__,
                    server.Stream.__name__,
                    request_type,
                    reply_type,
                )
                with buf.indent():
                    buf.add("pass")
            buf.add("")
            buf.add("def __mapping__(self) -> typing.Dict[str, {}.{}]:", const.__name__, const.Handler.__name__)
            with buf.indent():
                buf.add("return {{")
                with buf.indent():
                    for method in service.methods:
                        name, cardinality, request_type, reply_type = method
                        full_name = "/{}/{}".format(service_name, name)
                        buf.add("'{}': {}.{}(", full_name, const.__name__, const.Handler.__name__)
                        with buf.indent():
                            buf.add("self.{},", name)
                            buf.add("{}.{}.{},", const.__name__, const.Cardinality.__name__, cardinality.name)
                            buf.add("{},", request_type)
                            buf.add("{},", reply_type)
                        buf.add("),")
                buf.add("}}")

        buf.add("")
        buf.add("")
        buf.add("class Unimplemented{}Base({}Base):", service.name, service.name)
        with buf.indent():
            for name, _, request_type, reply_type in service.methods:
                buf.add("")
                buf.add(
                    "async def {}(self, stream: '{}.{}[{}, {}]') -> None:",
                    name,
                    server.__name__,
                    server.Stream.__name__,
                    request_type,
                    reply_type,
                )
                with buf.indent():
                    buf.add("raise grpclib.exceptions.GRPCError(grpclib.const.Status.UNIMPLEMENTED)")

        buf.add("")
        buf.add("")
        buf.add("class {}Stub:", service.name)
        with buf.indent():
            buf.add("")
            buf.add("def __init__(self, channel: {}.{}) -> None:".format(client.__name__, client.Channel.__name__))
            with buf.indent():
                if len(service.methods) == 0:
                    buf.add("pass")
                for method in service.methods:
                    name, cardinality, request_type, reply_type = method
                    full_name = "/{}/{}".format(service_name, name)
                    method_cls: type
                    if cardinality is const.Cardinality.UNARY_UNARY:
                        method_cls = client.UnaryUnaryMethod
                    elif cardinality is const.Cardinality.UNARY_STREAM:
                        method_cls = client.UnaryStreamMethod
                    elif cardinality is const.Cardinality.STREAM_UNARY:
                        method_cls = client.StreamUnaryMethod
                    elif cardinality is const.Cardinality.STREAM_STREAM:
                        method_cls = client.StreamStreamMethod
                    else:
                        raise TypeError(cardinality)
                    method_cls = cast(type, method_cls)  # FIXME: redundant
                    buf.add("self.{} = {}.{}(".format(name, client.__name__, method_cls.__name__))
                    with buf.indent():
                        buf.add("channel,")
                        buf.add("{!r},".format(full_name))
                        buf.add("{},", request_type)
                        buf.add("{},", reply_type)
                    buf.add(")")
    return buf.content()


def _get_proto(request: CodeGeneratorRequest, name: str) -> FileDescriptorProto:
    return next(f for f in request.proto_file if f.name == name)


def _strip_proto(proto_file_path: str) -> str:
    for suffix in [".protodevel", ".proto"]:
        if proto_file_path.endswith(suffix):
            return proto_file_path[: -len(suffix)]

    return proto_file_path


def _base_module_name(proto_file_path: str) -> str:
    basename = _strip_proto(proto_file_path)
    return basename.replace("-", "_").replace("/", ".")


def _proto2pb2_module_name(proto_file_path: str) -> str:
    return _base_module_name(proto_file_path) + "_pb2"


def _proto2grpc_module_name(proto_file_path: str) -> str:
    return _base_module_name(proto_file_path) + "_grpc"


def _type_names(
    proto_file: FileDescriptorProto,
    message_type: DescriptorProto,
    parents: Optional[Deque[str]] = None,
) -> Iterator[Tuple[str, str]]:
    if parents is None:
        parents = deque()

    proto_name_parts = [""]
    if proto_file.package:
        proto_name_parts.append(proto_file.package)
    proto_name_parts.extend(parents)
    proto_name_parts.append(message_type.name)

    py_name_parts = [_proto2pb2_module_name(proto_file.name)]
    py_name_parts.extend(parents)
    py_name_parts.append(message_type.name)

    yield ".".join(proto_name_parts), ".".join(py_name_parts)

    parents.append(message_type.name)
    for nested in message_type.nested_type:
        yield from _type_names(proto_file, nested, parents=parents)
    parents.pop()


def main() -> None:
    with os.fdopen(sys.stdin.fileno(), "rb") as inp:
        request = CodeGeneratorRequest.FromString(inp.read())

    types_map: Dict[str, str] = {}
    for pf in request.proto_file:
        for mt in pf.message_type:
            types_map.update(_type_names(pf, mt))

    response = CodeGeneratorResponse()

    # See https://github.com/protocolbuffers/protobuf/blob/v3.12.0/docs/implementing_proto3_presence.md  # noqa
    if hasattr(CodeGeneratorResponse, "Feature"):
        response.supported_features = CodeGeneratorResponse.FEATURE_PROTO3_OPTIONAL

    for file_to_generate in request.file_to_generate:
        proto_file = _get_proto(request, file_to_generate)

        imports = [_proto2pb2_module_name(dep) for dep in list(proto_file.dependency) + [file_to_generate]]

        services = []
        for service in proto_file.service:
            methods = []
            for method in service.method:
                cardinality = _CARDINALITY[(method.client_streaming, method.server_streaming)]  # type: ignore
                methods.append(
                    Method(
                        name=method.name,
                        cardinality=cardinality,
                        request_type=types_map[method.input_type],
                        reply_type=types_map[method.output_type],
                    )
                )
            services.append(Service(name=service.name, methods=methods))

        file = response.file.add()
        module_name = _proto2grpc_module_name(file_to_generate)
        file.name = module_name.replace(".", "/") + ".py"
        file.content = render(
            proto_file=proto_file.name,
            package=proto_file.package,
            imports=imports,
            services=services,
        )

    with os.fdopen(sys.stdout.fileno(), "wb") as out:
        out.write(response.SerializeToString())


if __name__ == "__main__":
    sys.argv[0] = re.sub(r"(-script\.pyw|\.exe)?$", "", sys.argv[0])
    sys.exit(main())