#!/usr/bin/env python3 # Copyright 2022 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # https://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Custom mmi2grpc gRPC compiler.""" from __future__ import annotations import os import sys from typing import Dict, List, Optional, Set, Tuple, Union from google.protobuf.compiler.plugin_pb2 import CodeGeneratorRequest, CodeGeneratorResponse from google.protobuf.descriptor import ( FieldDescriptor ) from google.protobuf.descriptor_pb2 import ( FileDescriptorProto, EnumDescriptorProto, DescriptorProto, ServiceDescriptorProto, MethodDescriptorProto, FieldDescriptorProto, ) _REQUEST = CodeGeneratorRequest.FromString(sys.stdin.buffer.read()) def find_type_in_file(proto_file: FileDescriptorProto, type_name: str) -> Optional[Union[DescriptorProto, EnumDescriptorProto]]: for enum in proto_file.enum_type: if enum.name == type_name: return enum for message in proto_file.message_type: if message.name == type_name: return message return None def find_type(package: str, type_name: str) -> Tuple[FileDescriptorProto, Union[DescriptorProto, EnumDescriptorProto]]: for file in _REQUEST.proto_file: if file.package == package and (type := find_type_in_file(file, type_name)): return file, type raise Exception(f'Type {package}.{type_name} not found') def add_import(imports: List[str], import_str: str) -> None: if not import_str in imports: imports.append(import_str) def import_type(imports: List[str], type: str, local: Optional[FileDescriptorProto]) -> Tuple[str, Union[DescriptorProto, EnumDescriptorProto], str]: package = type[1:type.rindex('.')] type_name = type[type.rindex('.')+1:] file, desc = find_type(package, type_name) if file == local: return f'{type_name}', desc, '' python_path = file.name.replace('.proto', '').replace('/', '.') module_path = python_path[:python_path.rindex('.')] module_name = python_path[python_path.rindex('.')+1:] + '_pb2' add_import(imports, f'from {module_path} import {module_name}') dft_import = '' if isinstance(desc, EnumDescriptorProto): dft_import = f'from {module_path}.{module_name} import {desc.value[0].name}' return f'{module_name}.{type_name}', desc, dft_import def collect_type(imports: List[str], parent: DescriptorProto, field: FieldDescriptorProto, local: Optional[FileDescriptorProto]) -> Tuple[str, str, str]: dft: str dft_import: str = '' if field.type == FieldDescriptor.TYPE_BYTES: type = 'bytes' dft = 'b\'\'' elif field.type == FieldDescriptor.TYPE_STRING: type = 'str' dft = '\'\'' elif field.type == FieldDescriptor.TYPE_BOOL: type = 'bool' dft = 'False' elif field.type in [ FieldDescriptor.TYPE_FLOAT, FieldDescriptor.TYPE_DOUBLE ]: type = 'float' dft = '0.0' elif field.type in [ FieldDescriptor.TYPE_INT64, FieldDescriptor.TYPE_UINT64, FieldDescriptor.TYPE_INT32, FieldDescriptor.TYPE_FIXED64, FieldDescriptor.TYPE_FIXED32, FieldDescriptor.TYPE_UINT32, FieldDescriptor.TYPE_SFIXED32, FieldDescriptor.TYPE_SFIXED64, FieldDescriptor.TYPE_SINT32, FieldDescriptor.TYPE_SINT64 ]: type = 'int' dft = '0' elif field.type in [FieldDescriptor.TYPE_ENUM, FieldDescriptor.TYPE_MESSAGE]: parts = field.type_name.split(f".{parent.name}.", 2) if len(parts) == 2: type = parts[1] for nested_type in parent.nested_type: if nested_type.name == type: assert nested_type.options.map_entry assert field.label == FieldDescriptor.LABEL_REPEATED key_type, _, _ = collect_type(imports, nested_type, nested_type.field[0], local) val_type, _, _ = collect_type(imports, nested_type, nested_type.field[1], local) add_import(imports, 'from typing import Dict') return f'Dict[{key_type}, {val_type}]', '{}', '' type, desc, enum_dft = import_type(imports, field.type_name, local) if isinstance(desc, EnumDescriptorProto): dft_import = enum_dft dft = desc.value[0].name else: dft = f'{type}()' else: raise Exception(f'TODO: {field}') if field.label == FieldDescriptor.LABEL_REPEATED: add_import(imports, 'from typing import List') type = f'List[{type}]' dft = '[]' return type, dft, dft_import def collect_field(imports: List[str], message: DescriptorProto, field: FieldDescriptorProto, local: Optional[FileDescriptorProto]) -> Tuple[Optional[int], str, str, str, str]: type, dft, dft_import = collect_type(imports, message, field, local) oneof_index = field.oneof_index if 'oneof_index' in f'{field}' else None return oneof_index, field.name, type, dft, dft_import def collect_message(imports: List[str], message: DescriptorProto, local: Optional[FileDescriptorProto]) -> Tuple[ List[Tuple[str, str, str]], Dict[str, List[Tuple[str, str]]], ]: fields: List[Tuple[str, str, str]] = [] oneof: Dict[str, List[Tuple[str, str]]] = {} for field in message.field: idx, name, type, dft, dft_import = collect_field(imports, message, field, local) if idx is not None: oneof_name = message.oneof_decl[idx].name oneof.setdefault(oneof_name, []) oneof[oneof_name].append((name, type)) else: add_import(imports, dft_import) fields.append((name, type, dft)) for oneof_name, oneof_fields in oneof.items(): for name, type in oneof_fields: add_import(imports, 'from typing import Optional') fields.append((name, f'Optional[{type}]', 'None')) return fields, oneof def generate_enum(imports: List[str], file: FileDescriptorProto, enum: EnumDescriptorProto, res: List[CodeGeneratorResponse.File]) -> List[str]: res.append(CodeGeneratorResponse.File( name=file.name.replace('.proto', '_pb2.py'), insertion_point=f'module_scope', content=f'class {enum.name}: ...\n\n' )) add_import(imports, 'from google.protobuf.internal.enum_type_wrapper import EnumTypeWrapper') return [ f'class {enum.name}(int, EnumTypeWrapper):', f' pass', f'', *[f'{value.name}: {enum.name}' for value in enum.value], '' ] def generate_message(imports: List[str], file: FileDescriptorProto, message: DescriptorProto, res: List[CodeGeneratorResponse.File]) -> List[str]: nested_message_lines: List[str] = [] message_lines: List[str] = [f'class {message.name}(Message):'] add_import(imports, 'from google.protobuf.message import Message') fields, oneof = collect_message(imports, message, file) for (name, type, _) in fields: message_lines.append(f' {name}: {type}') args = ', '.join([f'{name}: {type} = {dft}' for name, type, dft in fields]) if args: args = ', ' + args message_lines.extend([ f'', f' def __init__(self{args}) -> None: ...', f'' ]) for oneof_name, oneof_fields in oneof.items(): literals: str = ', '.join((f'Literal[\'{name}\']' for name, _ in oneof_fields)) types: Set[str] = set((type for _, type in oneof_fields)) if len(types) == 1: type = 'Optional[' + types.pop() + ']' else: types.add('None') type = 'Union[' + ', '.join(sorted(types)) + ']' nested_message_lines.extend([ f'class {message.name}_{oneof_name}_dict(TypedDict, total=False):', '\n'.join([f' {name}: {type}' for name, type in oneof_fields]), f'', ]) add_import(imports, 'from typing import Union') add_import(imports, 'from typing_extensions import TypedDict') add_import(imports, 'from typing_extensions import Literal') message_lines.extend([ f' @property', f' def {oneof_name}(self) -> {type}: ...' f'', f' def {oneof_name}_variant(self) -> Union[{literals}, None]: ...' f'', f' def {oneof_name}_asdict(self) -> {message.name}_{oneof_name}_dict: ...', f'', ]) return_variant = '\n '.join([f'if variant == \'{name}\': return unwrap(self.{name})' for name, _ in oneof_fields]) return_asdict = '\n '.join([f'if variant == \'{name}\': return {{\'{name}\': unwrap(self.{name})}} # type: ignore' for name, _ in oneof_fields]) if return_variant: return_variant += '\n ' if return_asdict: return_asdict += '\n ' res.append(CodeGeneratorResponse.File( name=file.name.replace('.proto', '_pb2.py'), insertion_point=f'module_scope', content=f""" def _{message.name}_{oneof_name}(self: {message.name}): variant = self.{oneof_name}_variant() if variant is None: return None {return_variant}raise Exception('Field `{oneof_name}` not found.') def _{message.name}_{oneof_name}_variant(self: {message.name}): return self.WhichOneof('{oneof_name}') # type: ignore def _{message.name}_{oneof_name}_asdict(self: {message.name}): variant = self.{oneof_name}_variant() if variant is None: return {{}} {return_asdict}raise Exception('Field `{oneof_name}` not found.') setattr({message.name}, '{oneof_name}', property(_{message.name}_{oneof_name})) setattr({message.name}, '{oneof_name}_variant', _{message.name}_{oneof_name}_variant) setattr({message.name}, '{oneof_name}_asdict', _{message.name}_{oneof_name}_asdict) """)) return message_lines + nested_message_lines def generate_service_method(imports: List[str], file: FileDescriptorProto, service: ServiceDescriptorProto, method: MethodDescriptorProto, sync: bool = True) -> List[str]: input_mode = 'stream' if method.client_streaming else 'unary' output_mode = 'stream' if method.server_streaming else 'unary' input_type, input_msg, _ = import_type(imports, method.input_type, None) output_type, _, _ = import_type(imports, method.output_type, None) input_type_pb2, _, _ = import_type(imports, method.input_type, None) output_type_pb2, _, _ = import_type(imports, method.output_type, None) if output_mode == 'stream': if input_mode == 'stream': output_type_hint = f'StreamStream[{input_type}, {output_type}]' if sync: add_import(imports, f'from ._utils import Sender') add_import(imports, f'from ._utils import Stream') add_import(imports, f'from ._utils import StreamStream') else: add_import(imports, f'from ._utils import AioSender as Sender') add_import(imports, f'from ._utils import AioStream as Stream') add_import(imports, f'from ._utils import AioStreamStream as StreamStream') else: output_type_hint = f'Stream[{output_type}]' if sync: add_import(imports, f'from ._utils import Stream') else: add_import(imports, f'from ._utils import AioStream as Stream') else: output_type_hint = output_type if sync else f'Awaitable[{output_type}]' if not sync: add_import(imports, f'from typing import Awaitable') if input_mode == 'stream' and output_mode == 'stream': add_import(imports, f'from typing import Optional') return ( f'def {method.name}(self, timeout: Optional[float] = None) -> {output_type_hint}:\n' f' tx: Sender[{input_type}] = Sender()\n' f' rx: Stream[{output_type}] = self.channel.{input_mode}_{output_mode}( # type: ignore\n' f" '/{file.package}.{service.name}/{method.name}',\n" f' request_serializer={input_type_pb2}.SerializeToString, # type: ignore\n' f' response_deserializer={output_type_pb2}.FromString # type: ignore\n' f' )(tx)\n' f' return StreamStream(tx, rx)' ).split('\n') if input_mode == 'stream': iterator_type = 'Iterator' if sync else 'AsyncIterator' add_import(imports, f'from typing import {iterator_type}') add_import(imports, f'from typing import Optional') return ( f'def {method.name}(self, iterator: {iterator_type}[{input_type}], timeout: Optional[float] = None) -> {output_type_hint}:\n' f' return self.channel.{input_mode}_{output_mode}( # type: ignore\n' f" '/{file.package}.{service.name}/{method.name}',\n" f' request_serializer={input_type_pb2}.SerializeToString, # type: ignore\n' f' response_deserializer={output_type_pb2}.FromString # type: ignore\n' f' )(iterator)' ).split('\n') else: add_import(imports, f'from typing import Optional') assert isinstance(input_msg, DescriptorProto) input_fields, _ = collect_message(imports, input_msg, None) args = ', '.join([f'{name}: {type} = {dft}' for name, type, dft in input_fields]) args_name = ', '.join([f'{name}={name}' for name, _, _ in input_fields]) if args: args = ', ' + args return ( f'def {method.name}(self{args}, wait_for_ready: Optional[bool] = None, timeout: Optional[float] = None) -> {output_type_hint}:\n' f' return self.channel.{input_mode}_{output_mode}( # type: ignore\n' f" '/{file.package}.{service.name}/{method.name}',\n" f' request_serializer={input_type_pb2}.SerializeToString, # type: ignore\n' f' response_deserializer={output_type_pb2}.FromString # type: ignore\n' f' )({input_type_pb2}({args_name}), wait_for_ready=wait_for_ready, timeout=timeout) # type: ignore' ).split('\n') def generate_service(imports: List[str], file: FileDescriptorProto, service: ServiceDescriptorProto, sync: bool = True) -> List[str]: methods = '\n\n '.join([ '\n '.join( generate_service_method(imports, file, service, method, sync) ) for method in service.method ]) channel_type = 'grpc.Channel' if sync else 'grpc.aio.Channel' return ( f'class {service.name}:\n' f' channel: {channel_type}\n' f'\n' f' def __init__(self, channel: {channel_type}) -> None:\n' f' self.channel = channel\n' f'\n' f' {methods}\n' ).split('\n') def generate_servicer_method(imports: List[str], method: MethodDescriptorProto, sync: bool = True) -> List[str]: input_mode = 'stream' if method.client_streaming else 'unary' output_mode = 'stream' if method.server_streaming else 'unary' input_type, _, _ = import_type(imports, method.input_type, None) output_type, _, _ = import_type(imports, method.output_type, None) output_type_hint = output_type if output_mode == 'stream': if sync: output_type_hint = f'Generator[{output_type}, None, None]' add_import(imports, f'from typing import Generator') else: output_type_hint = f'AsyncGenerator[{output_type}, None]' add_import(imports, f'from typing import AsyncGenerator') iterator_type = 'Iterator' if sync else 'AsyncIterator' if input_mode == 'stream': iterator_type = 'Iterator' if sync else 'AsyncIterator' add_import(imports, f'from typing import {iterator_type}') lines = (('' if sync else 'async ') + ( f'def {method.name}(self, request: {iterator_type}[{input_type}], context: grpc.ServicerContext) -> {output_type_hint}:\n' f' context.set_code(grpc.StatusCode.UNIMPLEMENTED) # type: ignore\n' f' context.set_details("Method not implemented!") # type: ignore\n' f' raise NotImplementedError("Method not implemented!")' )).split('\n') else: lines = (('' if sync else 'async ') + ( f'def {method.name}(self, request: {input_type}, context: grpc.ServicerContext) -> {output_type_hint}:\n' f' context.set_code(grpc.StatusCode.UNIMPLEMENTED) # type: ignore\n' f' context.set_details("Method not implemented!") # type: ignore\n' f' raise NotImplementedError("Method not implemented!")' )).split('\n') if output_mode == 'stream': lines.append(f' yield {output_type}() # no-op: to make the linter happy') return lines def generate_servicer(imports: List[str], file: FileDescriptorProto, service: ServiceDescriptorProto, sync: bool = True) -> List[str]: methods = '\n\n '.join([ '\n '.join( generate_servicer_method(imports, method, sync) ) for method in service.method ]) if not methods: methods = 'pass' return ( f'class {service.name}Servicer:\n' f' {methods}\n' ).split('\n') def generate_rpc_method_handler(imports: List[str], method: MethodDescriptorProto) -> List[str]: input_mode = 'stream' if method.client_streaming else 'unary' output_mode = 'stream' if method.server_streaming else 'unary' input_type, _, _ = import_type(imports, method.input_type, None) output_type, _, _ = import_type(imports, method.output_type, None) return ( f"'{method.name}': grpc.{input_mode}_{output_mode}_rpc_method_handler( # type: ignore\n" f' servicer.{method.name},\n' f' request_deserializer={input_type}.FromString, # type: ignore\n' f' response_serializer={output_type}.SerializeToString, # type: ignore\n' f' ),\n' ).split('\n') def generate_add_servicer_to_server_method(imports: List[str], file: FileDescriptorProto, service: ServiceDescriptorProto, sync: bool = True) -> List[str]: method_handlers = ' '.join([ '\n '.join( generate_rpc_method_handler(imports, method) ) for method in service.method ]) server_type = 'grpc.Server' if sync else 'grpc.aio.Server' return ( f'def add_{service.name}Servicer_to_server(servicer: {service.name}Servicer, server: {server_type}) -> None:\n' f' rpc_method_handlers = {{\n' f' {method_handlers}\n' f' }}\n' f' generic_handler = grpc.method_handlers_generic_handler( # type: ignore\n' f" '{file.package}.{service.name}', rpc_method_handlers)\n" f' server.add_generic_rpc_handlers((generic_handler,)) # type: ignore\n' ).split('\n') _HEADER = '''# Copyright 2022 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # https://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Generated python gRPC interfaces.""" from __future__ import annotations ''' _UTILS_PY = f'''{_HEADER} import asyncio import queue import grpc import sys from typing import Any, AsyncIterable, AsyncIterator, Generic, Iterator, TypeVar _T_co = TypeVar('_T_co', covariant=True) # pytype: disable=not-supported-yet _T = TypeVar('_T') class Stream(Iterator[_T_co], grpc.RpcContext): ... class AioStream(AsyncIterable[_T_co], grpc.RpcContext): ... class Sender(Iterator[_T]): if sys.version_info >= (3, 8): _inner: queue.Queue[_T] else: _inner: queue.Queue def __init__(self) -> None: self._inner = queue.Queue() def __iter__(self) -> Iterator[_T]: return self def __next__(self) -> _T: return self._inner.get() def send(self, item: _T) -> None: self._inner.put(item) class AioSender(AsyncIterator[_T]): if sys.version_info >= (3, 8): _inner: asyncio.Queue[_T] else: _inner: asyncio.Queue def __init__(self) -> None: self._inner = asyncio.Queue() def __iter__(self) -> AsyncIterator[_T]: return self async def __anext__(self) -> _T: return await self._inner.get() async def send(self, item: _T) -> None: await self._inner.put(item) def send_nowait(self, item: _T) -> None: self._inner.put_nowait(item) class StreamStream(Generic[_T, _T_co], Iterator[_T_co], grpc.RpcContext): _sender: Sender[_T] _receiver: Stream[_T_co] def __init__(self, sender: Sender[_T], receiver: Stream[_T_co]) -> None: self._sender = sender self._receiver = receiver def send(self, item: _T) -> None: self._sender.send(item) def __iter__(self) -> Iterator[_T_co]: return self._receiver.__iter__() def __next__(self) -> _T_co: return self._receiver.__next__() def is_active(self) -> bool: return self._receiver.is_active() # type: ignore def time_remaining(self) -> float: return self._receiver.time_remaining() # type: ignore def cancel(self) -> None: self._receiver.cancel() # type: ignore def add_callback(self, callback: Any) -> None: self._receiver.add_callback(callback) # type: ignore class AioStreamStream(Generic[_T, _T_co], AsyncIterator[_T_co], grpc.RpcContext): _sender: AioSender[_T] _receiver: AioStream[_T_co] def __init__(self, sender: AioSender[_T], receiver: AioStream[_T_co]) -> None: self._sender = sender self._receiver = receiver def __aiter__(self) -> AsyncIterator[_T_co]: return self._receiver.__aiter__() async def __anext__(self) -> _T_co: return await self._receiver.__aiter__().__anext__() async def send(self, item: _T) -> None: await self._sender.send(item) def send_nowait(self, item: _T) -> None: self._sender.send_nowait(item) def is_active(self) -> bool: return self._receiver.is_active() # type: ignore def time_remaining(self) -> float: return self._receiver.time_remaining() # type: ignore def cancel(self) -> None: self._receiver.cancel() # type: ignore def add_callback(self, callback: Any) -> None: self._receiver.add_callback(callback) # type: ignore ''' _FILES: List[CodeGeneratorResponse.File] = [] _UTILS_FILES: Set[str] = set() for file_name in _REQUEST.file_to_generate: file: FileDescriptorProto = next(filter(lambda x: x.name == file_name, _REQUEST.proto_file)) _FILES.append(CodeGeneratorResponse.File( name=file.name.replace('.proto', '_pb2.py'), insertion_point=f'module_scope', content='def unwrap(x):\n assert x\n return x\n' )) pyi_imports: List[str] = [] grpc_imports: List[str] = ['import grpc'] grpc_aio_imports: List[str] = ['import grpc', 'import grpc.aio'] enums = '\n'.join(sum([generate_enum(pyi_imports, file, enum, _FILES) for enum in file.enum_type], [])) messages = '\n'.join(sum([generate_message(pyi_imports, file, message, _FILES) for message in file.message_type], [])) services = '\n'.join(sum([generate_service(grpc_imports, file, service) for service in file.service], [])) aio_services = '\n'.join(sum([generate_service(grpc_aio_imports, file, service, False) for service in file.service], [])) servicers = '\n'.join(sum([generate_servicer(grpc_imports, file, service) for service in file.service], [])) aio_servicers = '\n'.join(sum([generate_servicer(grpc_aio_imports, file, service, False) for service in file.service], [])) add_servicer_methods = '\n'.join(sum([generate_add_servicer_to_server_method(grpc_imports, file, service) for service in file.service], [])) aio_add_servicer_methods = '\n'.join(sum([generate_add_servicer_to_server_method(grpc_aio_imports, file, service, False) for service in file.service], [])) pyi_imports.sort() grpc_imports.sort() grpc_aio_imports.sort() pyi_imports_str: str = '\n'.join(pyi_imports) grpc_imports_str: str = '\n'.join(grpc_imports) grpc_aio_imports_str: str = '\n'.join(grpc_aio_imports) utils_filename = file_name.replace(os.path.basename(file_name), '_utils.py') if utils_filename not in _UTILS_FILES: _UTILS_FILES.add(utils_filename) _FILES.extend([ CodeGeneratorResponse.File( name=utils_filename, content=_UTILS_PY, ) ]) _FILES.extend([ CodeGeneratorResponse.File( name=file.name.replace('.proto', '_pb2.pyi'), content=f'{_HEADER}\n\n{pyi_imports_str}\n\n{enums}\n\n{messages}\n' ), CodeGeneratorResponse.File( name=file_name.replace('.proto', '_grpc.py'), content=f'{_HEADER}\n\n{grpc_imports_str}\n\n{services}\n\n{servicers}\n\n{add_servicer_methods}' ), CodeGeneratorResponse.File( name=file_name.replace('.proto', '_grpc_aio.py'), content=f'{_HEADER}\n\n{grpc_aio_imports_str}\n\n{aio_services}\n\n{aio_servicers}\n\n{aio_add_servicer_methods}' ) ]) sys.stdout.buffer.write(CodeGeneratorResponse(file=_FILES).SerializeToString())