File tree Expand file tree Collapse file tree 3 files changed +21
-4
lines changed Expand file tree Collapse file tree 3 files changed +21
-4
lines changed Original file line number Diff line number Diff line change @@ -179,11 +179,14 @@ def comment(self) -> str:
179
179
proto_file = self .proto_file , path = self .path , indent = self .comment_indent
180
180
)
181
181
182
+ @dataclass
183
+ class Options :
184
+ grpc_kind : str = "grpclib"
182
185
183
186
@dataclass
184
187
class PluginRequestCompiler :
185
-
186
188
plugin_request_obj : plugin .CodeGeneratorRequest
189
+ options : Options
187
190
output_packages : Dict [str , "OutputTemplate" ] = field (default_factory = dict )
188
191
189
192
@property
Original file line number Diff line number Diff line change 29
29
FieldCompiler ,
30
30
MapEntryCompiler ,
31
31
MessageCompiler ,
32
+ Options ,
32
33
OneOfFieldCompiler ,
33
34
OutputTemplate ,
34
35
PluginRequestCompiler ,
41
42
if TYPE_CHECKING :
42
43
from google .protobuf .descriptor import Descriptor
43
44
44
-
45
45
def traverse (
46
46
proto_file : FieldDescriptorProto ,
47
47
) -> "itertools.chain[Tuple[Union[str, EnumDescriptorProto], List[int]]]" :
@@ -68,13 +68,21 @@ def _traverse(
68
68
_traverse ([5 ], proto_file .enum_type ), _traverse ([4 ], proto_file .message_type )
69
69
)
70
70
71
+ def parse_options (plugin_options : List [str ]) -> Options :
72
+ options = Options ()
73
+ for option in plugin_options :
74
+ if option .startswith ("grpc-kind=" ):
75
+ options .grpc_kind = option .split ("=" , 1 )[1 ]
76
+ return options
71
77
72
78
def generate_code (
73
79
request : plugin .CodeGeneratorRequest , response : plugin .CodeGeneratorResponse
74
80
) -> None :
75
81
plugin_options = request .parameter .split ("," ) if request .parameter else []
76
82
77
- request_data = PluginRequestCompiler (plugin_request_obj = request )
83
+ options = parse_options (plugin_options )
84
+
85
+ request_data = PluginRequestCompiler (plugin_request_obj = request , options = options )
78
86
# Gather output packages
79
87
for proto_file in request .proto_file :
80
88
if (
Original file line number Diff line number Diff line change @@ -15,8 +15,10 @@ from typing import {% for i in output_file.typing_imports|sort %}{{ i }}{% if no
15
15
{% endif %}
16
16
17
17
import betterproto
18
+ {% if output_file .parent_request .options .grpc_kind == "grpclib" %}
18
19
from betterproto.grpc.grpclib_server import ServiceBase
19
- {% if output_file .services %}
20
+ {% endif %}
21
+ {% if output_file .services and output_file .parent_request .options .grpc_kind == "grpclib" %}
20
22
import grpclib
21
23
{% endif %}
22
24
@@ -68,6 +70,9 @@ class {{ message.py_name }}(betterproto.Message):
68
70
69
71
70
72
{% endfor %}
73
+
74
+ {% if output_file .parent_request .options .grpc_kind == "grpclib" %}
75
+
71
76
{% for service in output_file .services %}
72
77
class {{ service.py_name }}Stub(betterproto.ServiceStub):
73
78
{% if service .comment %}
@@ -237,6 +242,7 @@ class {{ service.py_name }}Base(ServiceBase):
237
242
}
238
243
239
244
{% endfor %}
245
+ {% endif %}
240
246
241
247
{% for i in output_file .imports |sort %}
242
248
{{ i }}
You can’t perform that action at this time.
0 commit comments