1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
|
#!/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_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_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')
files = []
for file_name in request.file_to_generate:
file = next(filter(lambda x: x.name == file_name, request.proto_file))
imports = set([])
services = '\n'.join(sum([
generate_service(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
))
reponse = CodeGeneratorResponse(file=files)
sys.stdout.buffer.write(reponse.SerializeToString())
|
