service implementation stub generation

src/betterproto/__init__.py

@@ -25,6 +25,7 @@
 from ._types import T
 from .casing import camel_case, safe_snake_case, snake_case
 from .grpc.grpclib_client import ServiceStub
+from .grpc.grpclib_server import ServiceImplementation
 
 if sys.version_info[:2] < (3, 7):
     # Apply backport of datetime.fromisoformat from 3.7

src/betterproto/grpc/grpclib_server.py

@@ -0,0 +1,29 @@
+from abc import ABC
+
+import grpclib
+import grpclib.server
+
+
+class ServiceImplementation(ABC):
+    """
+    Base class for async gRPC servers.
+    """
+
+    __service_name__: str
+
+    def __rpc_methods__(self):
+        pass
+
+    def __mapping__(self):
+        mapping = {}
+        for (
+            method,
+            proto_name,
+            cardinality,
+            request_type,
+            response_type,
+        ) in self.__rpc_methods__():
+            mapping[f"/{self.__service_name__}/{proto_name}"] = grpclib.const.Handler(
+                method, cardinality, request_type, response_type
+            )
+        return mapping

src/betterproto/plugin/models.py

@@ -279,6 +279,10 @@ def proto_name(self) -> str:
     def py_name(self) -> str:
         return pythonize_class_name(self.proto_name)
 
+    @property
+    def py_name_as_field(self) -> str:
+        return pythonize_field_name(self.proto_name)
+
     @property
     def annotation(self) -> str:
         if self.repeated:
@@ -559,6 +563,10 @@ def __post_init__(self) -> None:
     def proto_name(self) -> str:
         return self.proto_obj.name
 
+    @property
+    def full_proto_name(self) -> str:
+        return f"{self.parent.package_proto_obj.package}.{self.proto_obj.name}"
+
     @property
     def py_name(self) -> str:
         return pythonize_class_name(self.proto_name)

src/betterproto/templates/template.py.j2

@@ -154,6 +154,97 @@ class {{ service.py_name }}Stub(betterproto.ServiceStub):
     {% endfor %}
 {% endfor %}
 
+{% for service in output_file.services %}
+class {{ service.py_name }}Implementation(betterproto.ServiceImplementation):
+    {% if service.comment %}
+{{ service.comment }}
+
+    {% endif %}
+
+    __service_name__ = "{{ service.full_proto_name }}"
+
+    {% for method in service.methods %}
+    async def {{ method.py_name }}(self
+        {%- if not method.client_streaming -%}
+            {%- if method.py_input_message and method.py_input_message.fields -%},
+                {%- for field in method.py_input_message.fields -%}
+                    {{ field.py_name }}: {% if field.py_name in method.mutable_default_args and not field.annotation.startswith("Optional[") -%}
+                                            Optional[{{ field.annotation }}]
+                                         {%- else -%}
+                                            {{ field.annotation }}
+                                         {%- endif -%}
+                    {%- if not loop.last %}, {% endif -%}
+                {%- endfor -%}
+            {%- endif -%}
+        {%- else -%}
+            {# Client streaming: need a request iterator instead #}
+            , {{ method.py_input_message.py_name_as_field }}_iterator: AsyncIterable["{{ method.py_input_message_type }}"]
+        {%- endif -%}
+            ) -> {% if method.server_streaming %}AsyncIterator["{{ method.py_output_message_type }}"]{% else %}"{{ method.py_output_message_type }}"{% endif %}:
+        {% if method.comment %}
+{{ method.comment }}
+
+        {% endif %}
+        raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)
+
+    {% endfor %}
+
+    {% for method in service.methods %}
+    async def __rpc_{{ method.py_name }}(self, stream):
+        {% if not method.client_streaming %}
+        request = await stream.recv_message()
+
+        request_kwargs = {
+        {% for field in method.py_input_message.fields %}
+            "{{ field.py_name }}": request.{{ field.py_name }},
+        {% endfor %}
+        }
+
+        {% else %}
+        request_kwargs = {"{{ method.py_input_message.py_name_as_field }}_iterator": stream.__aiter__()}
+        {% endif %}
+
+        {% if not method.server_streaming %}
+        response = await self.{{ method.py_name }}(**request_kwargs)
+        await stream.send_message(response)
+        {% else %}
+        response_iter = self.{{ method.py_name }}(**request_kwargs)
+        {# check if response is actually an AsyncIterator #}
+        {# this might be false if the method just returns without #}
+        {# yielding at least once #}
+        {# in that case, we just interpret it as an empty iterator #}
+        if isinstance(response_iter, AsyncIterable):
+            async for response_message in response_iter:
+                await stream.send_message(response_message)
+        else:
+            response_iter.close()
+        {% endif %}
+
+    {% endfor %}
+
+    def __rpc_methods__(self):
+        return [
+        {% for method in service.methods %}
+            (
+                self.__rpc_{{ method.py_name }},
+                "{{ method.proto_name }}",
+                {% if not method.client_streaming and not method.server_streaming %}
+                grpclib.const.Cardinality.UNARY_UNARY,
+                {% elif not method.client_streaming and method.server_streaming %}
+                grpclib.const.Cardinality.UNARY_STREAM,
+                {% elif method.client_streaming and not method.server_streaming %}
+                grpclib.const.Cardinality.STREAM_UNARY,
+                {% else %}
+                grpclib.const.Cardinality.STREAM_STREAM,
+                {% endif %}
+                {{ method.py_input_message_type }},
+                {{ method.py_output_message_type }},
+            ),
+        {% endfor %}
+    ]
+
+{% endfor %}
+
 {% for i in output_file.imports|sort %}
 {{ i }}
 {% endfor %}

tests/inputs/example_service/example_service.proto

@@ -0,0 +1,23 @@
+syntax = "proto3";
+
+package example_service;
+
+service ExampleService {
+  rpc ExampleUnaryUnary(ExampleRequest) returns (ExampleResponse);
+  rpc ExampleUnaryStream(ExampleRequest) returns (stream ExampleResponse);
+  rpc ExampleStreamUnary(stream ExampleRequest) returns (ExampleResponse);
+  rpc ExampleStreamStream(stream ExampleRequest) returns (stream ExampleResponse);
+}
+
+message ExampleRequest {
+  string example_string = 1;
+  int64  example_integer = 2;
+}
+
+message ExampleResponse {
+  string example_string = 1;
+  int64  example_integer = 2;
+}
+
+// Suppress test framework error when it's looking for a "Test" message or service
+message Test {}

tests/inputs/example_service/test_example_service.py

@@ -0,0 +1,108 @@
+import asyncio
+from typing import AsyncIterator, AsyncIterable
+
+from grpclib.client import Channel
+from grpclib.server import Server
+
+from tests.output_betterproto.example_service.example_service import (
+    ExampleServiceImplementation,
+    ExampleServiceStub,
+    ExampleRequest,
+    ExampleResponse,
+)
+
+
+class ExampleService(ExampleServiceImplementation):
+    async def example_unary_unary(
+        self, example_string: str, example_integer: int
+    ) -> "ExampleResponse":
+        return ExampleResponse(
+            example_string=example_string,
+            example_integer=example_integer,
+        )
+
+    async def example_unary_stream(
+        self, example_string: str, example_integer: int
+    ) -> AsyncIterator["ExampleResponse"]:
+        response = ExampleResponse(
+            example_string=example_string,
+            example_integer=example_integer,
+        )
+        yield response
+        yield response
+        yield response
+
+    async def example_stream_unary(
+        self, example_request_iterator: AsyncIterable["ExampleRequest"]
+    ) -> "ExampleResponse":
+        async for example_request in example_request_iterator:
+            return ExampleResponse(
+                example_string=example_request.example_string,
+                example_integer=example_request.example_integer,
+            )
+
+    async def example_stream_stream(
+        self, example_request_iterator: AsyncIterable["ExampleRequest"]
+    ) -> AsyncIterator["ExampleResponse"]:
+        async for example_request in example_request_iterator:
+            yield ExampleResponse(
+                example_string=example_request.example_string,
+                example_integer=example_request.example_integer,
+            )
+
+
+async def async_test_server_start():
+    host = "127.0.0.1"
+    port = 13337
+
+    test_string = "test string"
+    test_int = 42
+
+    # start server
+    server = Server([ExampleService()])
+    await server.start(host, port)
+
+    # start client
+    channel = Channel(host=host, port=port)
+    stub = ExampleServiceStub(channel)
+
+    # unary unary
+    response = await stub.example_unary_unary(
+        example_string="test string",
+        example_integer=42,
+    )
+    assert response.example_string == test_string
+    assert response.example_integer == test_int
+
+    # unary stream
+    async for response in stub.example_unary_stream(
+        example_string="test string",
+        example_integer=42,
+    ):
+        assert response.example_string == test_string
+        assert response.example_integer == test_int
+
+    # stream unary
+    request = ExampleRequest(
+        example_string=test_string,
+        example_integer=42,
+    )
+
+    async def request_iterator():
+        yield request
+        yield request
+        yield request
+
+    response = await stub.example_stream_unary(request_iterator())
+    assert response.example_string == test_string
+    assert response.example_integer == test_int
+
+    # stream stream
+    async for response in stub.example_stream_stream(request_iterator()):
+        assert response.example_string == test_string
+        assert response.example_integer == test_int
+
+
+def test_server_start():
+    loop = asyncio.get_event_loop()
+    loop.run_until_complete(async_test_server_start())