Skip to content

Commit 0fa1754

Browse files
committed
Add grpcio base generation support
1 parent 614a5f2 commit 0fa1754

File tree

2 files changed

+61
-0
lines changed

2 files changed

+61
-0
lines changed

src/betterproto/plugin/models.py

+4
Original file line numberDiff line numberDiff line change
@@ -563,6 +563,10 @@ def __post_init__(self) -> None:
563563
def proto_name(self) -> str:
564564
return self.proto_obj.name
565565

566+
@property
567+
def proto_path(self) -> str:
568+
return self.parent.package + "." + self.proto_name
569+
566570
@property
567571
def py_name(self) -> str:
568572
return pythonize_class_name(self.proto_name)

src/betterproto/templates/template.py.j2

+57
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,9 @@ from betterproto.grpc.grpclib_server import ServiceBase
2121
{% if output_file.services and output_file.parent_request.options.grpc_kind == "grpclib" %}
2222
import grpclib
2323
{% endif %}
24+
{% if output_file.services and output_file.parent_request.options.grpc_kind == "grpcio" %}
25+
import grpc
26+
{% endif %}
2427

2528

2629
{% if output_file.enums %}{% for enum in output_file.enums %}
@@ -244,6 +247,60 @@ class {{ service.py_name }}Base(ServiceBase):
244247
{% endfor %}
245248
{% endif %}
246249

250+
{% if output_file.parent_request.options.grpc_kind == "grpcio" %}
251+
{% for service in output_file.services %}
252+
class {{ service.py_name }}Base:
253+
{% if service.comment %}
254+
{{ service.comment }}
255+
256+
{% endif %}
257+
258+
{% for method in service.methods %}
259+
async def {{ method.py_name }}(self
260+
{%- if not method.client_streaming -%}
261+
, request: "{{ method.py_input_message_type }}"
262+
{%- else -%}
263+
{# Client streaming: need a request iterator instead #}
264+
, request_iterator: AsyncIterator["{{ method.py_input_message_type }}"]
265+
{%- endif -%}
266+
, context: grpc.aio.ServicerContext
267+
) -> {% if method.server_streaming %}AsyncGenerator["{{ method.py_output_message_type }}", None]{% else %}"{{ method.py_output_message_type }}"{% endif %}:
268+
{% if method.comment %}
269+
{{ method.comment }}
270+
271+
{% endif %}
272+
context.set_code(grpc.StatusCode.UNIMPLEMENTED)
273+
context.set_details('Method not implemented!')
274+
raise NotImplementedError('Method not implemented!')
275+
276+
{% endfor %}
277+
278+
def register_to_server(self, server: grpc.aio.Server):
279+
rpc_method_handlers = {
280+
{% for method in service.methods %}
281+
"{{ method.proto_name }}":
282+
{% if not method.client_streaming and not method.server_streaming %}
283+
grpc.unary_unary_rpc_method_handler(
284+
{% elif method.client_streaming and method.server_streaming %}
285+
grpc.stream_stream_rpc_method_handler(
286+
{% elif method.client_streaming %}
287+
grpc.stream_unary_rpc_method_handler(
288+
{% else %}
289+
grpc.unary_stream_rpc_method_handler(
290+
{% endif %}
291+
self.{{ method.py_name }},
292+
request_deserializer={{ method.py_input_message_type }}.FromString,
293+
response_serializer={{ method.py_input_message_type }}.SerializeToString,
294+
),
295+
{% endfor %}
296+
}
297+
generic_handler = grpc.method_handlers_generic_handler(
298+
"{{ service.proto_path }}", rpc_method_handlers)
299+
self.add_generic_rpc_handlers((generic_handler,))
300+
301+
{% endfor %}
302+
{% endif %}
303+
247304
{% for i in output_file.imports|sort %}
248305
{{ i }}
249306
{% endfor %}

0 commit comments

Comments
 (0)