Add grpcio support.

Author: mokucher

README.md

@@ -533,6 +533,12 @@ protoc \
     /usr/local/include/google/protobuf/*.proto
 ```
 
+### Using grpcio library instead of grpclib
+
+In order to use the `grpcio` library instead of `grpclib`, you can use the `--custom_opt=grpcio`
+option when running the `protoc` command. 
+This will generate stubs compatible with the `grpcio` library.
+
 ### TODO
 
 - [x] Fixed length fields

poetry.lock

@@ -19,6 +19,7 @@ dynamic = ["dependencies"]
 # The Ruff version is pinned. To update it, also update it in .pre-commit-config.yaml
 ruff = { version = "~0.9.1", optional = true }
 grpclib = "^0.4.1"
+grpcio = { version = ">=1.73.0", optional = true }
 jinja2 = { version = ">=3.0.3", optional = true }
 python-dateutil = "^2.8"
 typing-extensions = "^4.7.1"
@@ -45,13 +46,15 @@ pydantic = ">=2.0,<3"
 protobuf = "^5"
 cachelib = "^0.13.0"
 tomlkit = ">=0.7.0"
+grpcio-testing = "^1.54.2"
 
 [project.scripts]
 protoc-gen-python_betterproto = "betterproto.plugin:main"
 
 [project.optional-dependencies]
 compiler = ["ruff", "jinja2"]
 rust-codec = ["betterproto-rust-codec"]
+grpcio = ["grpcio"]
 
 [tool.ruff]
 extend-exclude = ["tests/output_*"]

pyproject.toml

@@ -0,0 +1,130 @@
+from abc import ABC
+from typing import (
+    TYPE_CHECKING,
+    AsyncIterable,
+    AsyncIterator,
+    Iterable,
+    Mapping,
+    Optional,
+    Union,
+)
+
+import grpc
+
+
+if TYPE_CHECKING:
+    from .._types import (
+        T,
+        IProtoMessage,
+    )
+
+
+Value = Union[str, bytes]
+MetadataLike = Union[Mapping[str, Value], Iterable[tuple[str, Value]]]
+MessageSource = Union[Iterable["IProtoMessage"], AsyncIterable["IProtoMessage"]]
+
+
+class ServiceStub(ABC):
+
+    def __init__(
+        self,
+        channel: grpc.aio.Channel,
+        *,
+        timeout: Optional[float] = None,
+        metadata: Optional[MetadataLike] = None,
+    ) -> None:
+        self.channel = channel
+        self.timeout = timeout
+        self.metadata = metadata
+
+    def _resolve_request_kwargs(
+        self,
+        timeout: Optional[float],
+        metadata: Optional[MetadataLike],
+    ):
+        # Avoid creating dict if no overrides needed
+        if timeout is None and metadata is None:
+            # Return cached default kwargs if both timeout and metadata are None
+            if not hasattr(self, '_default_kwargs'):
+                self._default_kwargs = {
+                    "timeout": self.timeout,
+                    "metadata": self.metadata,
+                }
+            return self._default_kwargs
+
+        return {
+            "timeout": self.timeout if timeout is None else timeout,
+            "metadata": self.metadata if metadata is None else metadata,
+        }
+
+    async def _unary_unary(
+        self,
+        stub_method: grpc.aio.UnaryUnaryMultiCallable,
+        request: "IProtoMessage",
+        *,
+        timeout: Optional[float] = None,
+        metadata: Optional[MetadataLike] = None,
+    ) -> "T":
+        return await stub_method(
+            request,
+            **self._resolve_request_kwargs(timeout, metadata),
+        )
+
+    async def _unary_stream(
+        self,
+        stub_method: grpc.aio.UnaryStreamMultiCallable,
+        request: "IProtoMessage",
+        *,
+        timeout: Optional[float] = None,
+        metadata: Optional[MetadataLike] = None,
+    ) -> AsyncIterator["T"]:
+        call = stub_method(
+            request,
+            **self._resolve_request_kwargs(timeout, metadata),
+        )
+        async for response in call:
+            yield response
+
+    async def _stream_unary(
+        self,
+        stub_method: grpc.aio.StreamUnaryMultiCallable,
+        request_iterator: MessageSource,
+        *,
+        timeout: Optional[float] = None,
+        metadata: Optional[MetadataLike] = None,
+    ) -> "T":
+        call = stub_method(
+            self._wrap_message_iterator(request_iterator),
+            **self._resolve_request_kwargs(timeout, metadata),
+        )
+        return await call
+
+    async def _stream_stream(
+        self,
+        stub_method: grpc.aio.StreamStreamMultiCallable,
+        request_iterator: MessageSource,
+        *,
+        timeout: Optional[float] = None,
+        metadata: Optional[MetadataLike] = None,
+    ) -> AsyncIterator["T"]:
+        call = stub_method(
+            self._wrap_message_iterator(request_iterator),
+            **self._resolve_request_kwargs(timeout, metadata),
+        )
+        async for response in call:
+            yield response
+
+    @staticmethod
+    def _wrap_message_iterator(
+        messages: MessageSource,
+    ) -> AsyncIterator["IProtoMessage"]:
+        if hasattr(messages, '__aiter__'):
+            async def async_wrapper():
+                async for message in messages:
+                    yield message
+            return async_wrapper()
+        else:
+            async def sync_wrapper():
+                for message in messages:
+                    yield message
+            return sync_wrapper()

src/betterproto/grpc/grpcio_client.py

@@ -0,0 +1,30 @@
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Dict
+
+
+if TYPE_CHECKING:
+    import grpc
+
+
+class ServiceBase(ABC):
+
+    @property
+    @abstractmethod
+    def __rpc_methods__(self) -> Dict[str, "grpc.RpcMethodHandler"]: ...
+
+    @property
+    @abstractmethod
+    def __proto_path__(self) -> str: ...
+
+
+def register_servicers(server: "grpc.aio.Server", *servicers: ServiceBase):
+    from grpc import method_handlers_generic_handler
+
+    server.add_generic_rpc_handlers(
+        tuple(
+            method_handlers_generic_handler(
+                servicer.__proto_path__, servicer.__rpc_methods__
+            )
+            for servicer in servicers
+        )
+    )

src/betterproto/grpc/grpcio_server.py

@@ -270,6 +270,7 @@ class OutputTemplate:
     imports_type_checking_only: Set[str] = field(default_factory=set)
     pydantic_dataclasses: bool = False
     output: bool = True
+    use_grpcio: bool = False
     typing_compiler: TypingCompiler = field(default_factory=DirectImportTypingCompiler)
 
     @property
@@ -697,18 +698,24 @@ class ServiceMethodCompiler(ProtoContentBase):
     proto_obj: MethodDescriptorProto
     path: List[int] = PLACEHOLDER
     comment_indent: int = 8
+    use_grpcio: bool = False
 
     def __post_init__(self) -> None:
         # Add method to service
         self.parent.methods.append(self)
-
-        self.output_file.imports_type_checking_only.add("import grpclib.server")
-        self.output_file.imports_type_checking_only.add(
-            "from betterproto.grpc.grpclib_client import MetadataLike"
-        )
-        self.output_file.imports_type_checking_only.add(
-            "from grpclib.metadata import Deadline"
-        )
+        if self.use_grpcio:
+            self.output_file.imports_type_checking_only.add("import grpc.aio")
+            self.output_file.imports_type_checking_only.add(
+                "from betterproto.grpc.grpcio_client import MetadataLike"
+            )
+        else:
+            self.output_file.imports_type_checking_only.add("import grpclib.server")
+            self.output_file.imports_type_checking_only.add(
+                "from betterproto.grpc.grpclib_client import MetadataLike"
+            )
+            self.output_file.imports_type_checking_only.add(
+                "from grpclib.metadata import Deadline"
+            )
 
         super().__post_init__()  # check for unset fields
 

src/betterproto/plugin/models.py

@@ -80,6 +80,7 @@ def generate_code(request: CodeGeneratorRequest) -> CodeGeneratorResponse:
     response.supported_features = CodeGeneratorResponseFeature.FEATURE_PROTO3_OPTIONAL
 
     request_data = PluginRequestCompiler(plugin_request_obj=request)
+    use_grpcio = "USE_GRPCIO" in plugin_options
     # Gather output packages
     for proto_file in request.proto_file:
         output_package_name = proto_file.package
@@ -90,7 +91,8 @@ def generate_code(request: CodeGeneratorRequest) -> CodeGeneratorResponse:
             )
         # Add this input file to the output corresponding to this package
         request_data.output_packages[output_package_name].input_files.append(proto_file)
-
+        if use_grpcio:
+            request_data.output_packages[output_package_name].use_grpcio = True
         if (
             proto_file.package == "google.protobuf"
             and "INCLUDE_GOOGLE" not in plugin_options
@@ -143,7 +145,9 @@ def generate_code(request: CodeGeneratorRequest) -> CodeGeneratorResponse:
     for output_package_name, output_package in request_data.output_packages.items():
         for proto_input_file in output_package.input_files:
             for index, service in enumerate(proto_input_file.service):
-                read_protobuf_service(proto_input_file, service, index, output_package)
+                read_protobuf_service(
+                    proto_input_file, service, index, output_package, use_grpcio
+                )
 
     # Generate output files
     output_paths: Set[pathlib.Path] = set()
@@ -253,6 +257,7 @@ def read_protobuf_service(
     service: ServiceDescriptorProto,
     index: int,
     output_package: OutputTemplate,
+    use_grpcio: bool = False,
 ) -> None:
     service_data = ServiceCompiler(
         source_file=source_file,
@@ -266,4 +271,5 @@ def read_protobuf_service(
             parent=service_data,
             proto_obj=method,
             path=[6, index, 2, j],
+            use_grpcio=use_grpcio,
         )

src/betterproto/plugin/parser.py

@@ -42,8 +42,16 @@ from pydantic import {% for i in output_file.pydantic_imports|sort %}{{ i }}{% i
 
 {% endif %}
 
+
+{% if output_file.use_grpcio %}
+import grpc
+from betterproto.grpc.grpcio_client import ServiceStub
+from betterproto.grpc.grpcio_server import ServiceBase
+{% endif %}
+
 import betterproto
-{% if output_file.services %}
+{% if not output_file.use_grpcio %}
+from betterproto.grpc.grpclib_client import ServiceStub
 from betterproto.grpc.grpclib_server import ServiceBase
 import grpclib
 {% endif %}