Skip to content

Commit 6957b07

Browse files
author
robin
committed
Add grpcio base generation support
1 parent 71428cc commit 6957b07

File tree

2 files changed

+61
-0
lines changed

2 files changed

+61
-0
lines changed

src/betterproto/plugin/models.py

Lines changed: 4 additions & 0 deletions
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

Lines changed: 57 additions & 0 deletions
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+
server.add_generic_rpc_handlers(grpc.method_handlers_generic_handler(
280+
"{{ service.proto_path }}",
281+
{
282+
{% for method in service.methods %}
283+
"{{ method.proto_name }}":
284+
{% if not method.client_streaming and not method.server_streaming %}
285+
grpc.unary_unary_rpc_method_handler(
286+
{% elif method.client_streaming and method.server_streaming %}
287+
grpc.stream_stream_rpc_method_handler(
288+
{% elif method.client_streaming %}
289+
grpc.stream_unary_rpc_method_handler(
290+
{% else %}
291+
grpc.unary_stream_rpc_method_handler(
292+
{% endif %}
293+
self.{{ method.py_name }},
294+
request_deserializer={{ method.py_input_message_type }}.FromString,
295+
response_serializer={{ method.py_input_message_type }}.SerializeToString,
296+
),
297+
{% endfor %}
298+
}
299+
))
300+
301+
{% endfor %}
302+
{% endif %}
303+
247304
{% for i in output_file.imports|sort %}
248305
{{ i }}
249306
{% endfor %}

0 commit comments

Comments
 (0)