184 lines
6.3 KiB
Python
Executable File
184 lines
6.3 KiB
Python
Executable File
#!/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."""
|
|
|
|
import sys
|
|
|
|
from google.protobuf.compiler.plugin_pb2 import CodeGeneratorRequest, \
|
|
CodeGeneratorResponse
|
|
|
|
|
|
def eprint(*args, **kwargs):
|
|
print(*args, file=sys.stderr, **kwargs)
|
|
|
|
|
|
request = CodeGeneratorRequest.FromString(sys.stdin.buffer.read())
|
|
|
|
|
|
def has_type(proto_file, type_name):
|
|
return any(filter(lambda x: x.name == type_name, proto_file.message_type))
|
|
|
|
|
|
def import_type(imports, type):
|
|
package = type[1:type.rindex('.')]
|
|
type_name = type[type.rindex('.')+1:]
|
|
file = next(filter(
|
|
lambda x: x.package == package and has_type(x, type_name),
|
|
request.proto_file))
|
|
python_path = file.name.replace('.proto', '').replace('/', '.')
|
|
as_name = python_path.replace('.', '_dot_') + '__pb2'
|
|
module_path = python_path[:python_path.rindex('.')]
|
|
module_name = python_path[python_path.rindex('.')+1:] + '_pb2'
|
|
imports.add(f'from {module_path} import {module_name} as {as_name}')
|
|
return f'{as_name}.{type_name}'
|
|
|
|
|
|
def generate_service_method(imports, file, service, method):
|
|
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)
|
|
output_type = import_type(imports, method.output_type)
|
|
|
|
if input_mode == 'stream':
|
|
return (
|
|
f'def {method.name}(self, iterator, **kwargs):\n'
|
|
f' return self.channel.{input_mode}_{output_mode}(\n'
|
|
f" '/{file.package}.{service.name}/{method.name}',\n"
|
|
f' request_serializer={input_type}.SerializeToString,\n'
|
|
f' response_deserializer={output_type}.FromString\n'
|
|
f' )(iterator, **kwargs)'
|
|
).split('\n')
|
|
else:
|
|
return (
|
|
f'def {method.name}(self, wait_for_ready=None, **kwargs):\n'
|
|
f' return self.channel.{input_mode}_{output_mode}(\n'
|
|
f" '/{file.package}.{service.name}/{method.name}',\n"
|
|
f' request_serializer={input_type}.SerializeToString,\n'
|
|
f' response_deserializer={output_type}.FromString\n'
|
|
f' )({input_type}(**kwargs), wait_for_ready=wait_for_ready)'
|
|
).split('\n')
|
|
|
|
|
|
def generate_service(imports, file, service):
|
|
methods = '\n\n '.join([
|
|
'\n '.join(
|
|
generate_service_method(imports, file, service, method)
|
|
) for method in service.method
|
|
])
|
|
return (
|
|
f'class {service.name}:\n'
|
|
f' def __init__(self, channel):\n'
|
|
f' self.channel = channel\n'
|
|
f'\n'
|
|
f' {methods}\n'
|
|
).split('\n')
|
|
|
|
|
|
def generate_servicer_method(method):
|
|
input_mode = 'stream' if method.client_streaming else 'unary'
|
|
|
|
if input_mode == 'stream':
|
|
return (
|
|
f'def {method.name}(self, request_iterator, context):\n'
|
|
f' context.set_code(grpc.StatusCode.UNIMPLEMENTED)\n'
|
|
f' context.set_details("Method not implemented!")\n'
|
|
f' raise NotImplementedError("Method not implemented!")'
|
|
).split('\n')
|
|
else:
|
|
return (
|
|
f'def {method.name}(self, request, context):\n'
|
|
f' context.set_code(grpc.StatusCode.UNIMPLEMENTED)\n'
|
|
f' context.set_details("Method not implemented!")\n'
|
|
f' raise NotImplementedError("Method not implemented!")'
|
|
).split('\n')
|
|
|
|
|
|
def generate_servicer(service):
|
|
methods = '\n\n '.join([
|
|
'\n '.join(
|
|
generate_servicer_method(method)
|
|
) for method in service.method
|
|
])
|
|
return (
|
|
f'class {service.name}Servicer:\n'
|
|
f'\n'
|
|
f' {methods}\n'
|
|
).split('\n')
|
|
|
|
|
|
def generate_rpc_method_handler(imports, method):
|
|
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)
|
|
output_type = import_type(imports, method.output_type)
|
|
|
|
return (
|
|
f"'{method.name}': grpc.{input_mode}_{output_mode}_rpc_method_handler(\n"
|
|
f' servicer.{method.name},\n'
|
|
f' request_deserializer={input_type}.FromString,\n'
|
|
f' response_serializer={output_type}.SerializeToString,\n'
|
|
f' ),\n'
|
|
).split('\n')
|
|
|
|
|
|
def generate_add_servicer_to_server_method(imports, file, service):
|
|
method_handlers = ' '.join([
|
|
'\n '.join(
|
|
generate_rpc_method_handler(imports, method)
|
|
) for method in service.method
|
|
])
|
|
return (
|
|
f'def add_{service.name}Servicer_to_server(servicer, server):\n'
|
|
f' rpc_method_handlers = {{\n'
|
|
f' {method_handlers}\n'
|
|
f' }}\n'
|
|
f' generic_handler = grpc.method_handlers_generic_handler(\n'
|
|
f" '{file.package}.{service.name}', rpc_method_handlers)\n"
|
|
f' server.add_generic_rpc_handlers((generic_handler,))'
|
|
).split('\n')
|
|
|
|
|
|
files = []
|
|
|
|
for file_name in request.file_to_generate:
|
|
file = next(filter(lambda x: x.name == file_name, request.proto_file))
|
|
|
|
imports = set(['import grpc'])
|
|
|
|
services = '\n'.join(sum([
|
|
generate_service(imports, file, service) for service in file.service
|
|
], []))
|
|
|
|
servicers = '\n'.join(sum([
|
|
generate_servicer(service) for service in file.service
|
|
], []))
|
|
|
|
add_servicer_methods = '\n'.join(sum([
|
|
generate_add_servicer_to_server_method(imports, file, service) for service in file.service
|
|
], []))
|
|
|
|
files.append(CodeGeneratorResponse.File(
|
|
name=file_name.replace('.proto', '_grpc.py'),
|
|
content='\n'.join(imports) + '\n\n' + services + '\n\n' + servicers + '\n\n' + add_servicer_methods + '\n'
|
|
))
|
|
|
|
response = CodeGeneratorResponse(file=files)
|
|
|
|
sys.stdout.buffer.write(response.SerializeToString())
|