pub/python-betterproto
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Expose timeout, deadline and metadata parameters from grpclib (#352)

Browse Source
This commit is contained in:
Arun Babu Neelicattu
2022-03-13 23:34:11 +01:00
committed by GitHub
parent 62da35b3ea
commit 18a518efa7
4 changed files with 101 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

30
src/betterproto/grpc/grpclib_client.py
View File

@@ -22,10 +22,10 @@ if TYPE_CHECKING:
from grpclib.metadata import Deadline
_Value = Union[str, bytes]
_MetadataLike = Union[Mapping[str, _Value], Collection[Tuple[str, _Value]]]
_MessageLike = Union[T, ST]
_MessageSource = Union[Iterable[ST], AsyncIterable[ST]]
Value = Union[str, bytes]
MetadataLike = Union[Mapping[str, Value], Collection[Tuple[str, Value]]]
MessageLike = Union[T, ST]
MessageSource = Union[Iterable[ST], AsyncIterable[ST]]
class ServiceStub(ABC):
@@ -39,7 +39,7 @@ class ServiceStub(ABC):
*,
timeout: Optional[float] = None,
deadline: Optional["Deadline"] = None,
metadata: Optional[_MetadataLike] = None,
metadata: Optional[MetadataLike] = None,
) -> None:
self.channel = channel
self.timeout = timeout
@@ -50,7 +50,7 @@ class ServiceStub(ABC):
self,
timeout: Optional[float],
deadline: Optional["Deadline"],
metadata: Optional[_MetadataLike],
metadata: Optional[MetadataLike],
):
return {
"timeout": self.timeout if timeout is None else timeout,
@@ -61,12 +61,12 @@ class ServiceStub(ABC):
async def _unary_unary(
self,
route: str,
request: _MessageLike,
request: MessageLike,
response_type: Type[T],
*,
timeout: Optional[float] = None,
deadline: Optional["Deadline"] = None,
metadata: Optional[_MetadataLike] = None,
metadata: Optional[MetadataLike] = None,
) -> T:
"""Make a unary request and return the response."""
async with self.channel.request(
@@ -84,12 +84,12 @@ class ServiceStub(ABC):
async def _unary_stream(
self,
route: str,
request: _MessageLike,
request: MessageLike,
response_type: Type[T],
*,
timeout: Optional[float] = None,
deadline: Optional["Deadline"] = None,
metadata: Optional[_MetadataLike] = None,
metadata: Optional[MetadataLike] = None,
) -> AsyncIterator[T]:
"""Make a unary request and return the stream response iterator."""
async with self.channel.request(
@@ -106,13 +106,13 @@ class ServiceStub(ABC):
async def _stream_unary(
self,
route: str,
request_iterator: _MessageSource,
request_iterator: MessageSource,
request_type: Type[ST],
response_type: Type[T],
*,
timeout: Optional[float] = None,
deadline: Optional["Deadline"] = None,
metadata: Optional[_MetadataLike] = None,
metadata: Optional[MetadataLike] = None,
) -> T:
"""Make a stream request and return the response."""
async with self.channel.request(
@@ -130,13 +130,13 @@ class ServiceStub(ABC):
async def _stream_stream(
self,
route: str,
request_iterator: _MessageSource,
request_iterator: MessageSource,
request_type: Type[ST],
response_type: Type[T],
*,
timeout: Optional[float] = None,
deadline: Optional["Deadline"] = None,
metadata: Optional[_MetadataLike] = None,
metadata: Optional[MetadataLike] = None,
) -> AsyncIterator[T]:
"""
Make a stream request and return an AsyncIterator to iterate over response
@@ -161,7 +161,7 @@ class ServiceStub(ABC):
raise
@staticmethod
async def _send_messages(stream, messages: _MessageSource):
async def _send_messages(stream, messages: MessageSource):
if isinstance(messages, AsyncIterable):
async for message in messages:
await stream.send_message(message)

10
src/betterproto/plugin/models.py
View File

@@ -232,6 +232,7 @@ class OutputTemplate:
messages: List["MessageCompiler"] = field(default_factory=list)
enums: List["EnumDefinitionCompiler"] = field(default_factory=list)
services: List["ServiceCompiler"] = field(default_factory=list)
imports_type_checking_only: Set[str] = field(default_factory=set)
@property
def package(self) -> str:
@@ -679,6 +680,15 @@ class ServiceMethodCompiler(ProtoContentBase):
if self.client_streaming or self.server_streaming:
self.output_file.typing_imports.add("AsyncIterator")
# add imports required for request arguments timeout, deadline and metadata
self.output_file.typing_imports.add("Optional")
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
@property

26
src/betterproto/templates/template.py.j2
View File

@@ -20,6 +20,13 @@ from betterproto.grpc.grpclib_server import ServiceBase
import grpclib
{% endif %}
{% if output_file.imports_type_checking_only %}
from typing import TYPE_CHECKING
if TYPE_CHECKING:
{% for i in output_file.imports_type_checking_only|sort %} {{ i }}
{% endfor %}
{% endif %}
{% if output_file.enums %}{% for enum in output_file.enums %}
class {{ enum.py_name }}(betterproto.Enum):
@@ -86,6 +93,9 @@ class {{ service.py_name }}Stub(betterproto.ServiceStub):
{# Client streaming: need a request iterator instead #}
, {{ method.py_input_message_param }}_iterator: Union[AsyncIterable["{{ method.py_input_message_type }}"], Iterable["{{ method.py_input_message_type }}"]]
{%- endif -%}
, timeout: Optional[float] = None
, deadline: Optional["Deadline"] = None
, metadata: Optional["_MetadataLike"] = None
) -> {% if method.server_streaming %}AsyncIterator["{{ method.py_output_message_type }}"]{% else %}"{{ method.py_output_message_type }}"{% endif %}:
{% if method.comment %}
{{ method.comment }}
@@ -98,6 +108,9 @@ class {{ service.py_name }}Stub(betterproto.ServiceStub):
{{ method.py_input_message_param }}_iterator,
{{ method.py_input_message_type }},
{{ method.py_output_message_type.strip('"') }},
timeout=timeout,
deadline=deadline,
metadata=metadata,
):
yield response
{% else %}{# i.e. not client streaming #}
@@ -105,6 +118,9 @@ class {{ service.py_name }}Stub(betterproto.ServiceStub):
"{{ method.route }}",
{{ method.py_input_message_param }},
{{ method.py_output_message_type.strip('"') }},
timeout=timeout,
deadline=deadline,
metadata=metadata,
):
yield response
@@ -115,13 +131,19 @@ class {{ service.py_name }}Stub(betterproto.ServiceStub):
"{{ method.route }}",
{{ method.py_input_message_param }}_iterator,
{{ method.py_input_message_type }},
{{ method.py_output_message_type.strip('"') }}
{{ method.py_output_message_type.strip('"') }},
timeout=timeout,
deadline=deadline,
metadata=metadata,
)
{% else %}{# i.e. not client streaming #}
return await self._unary_unary(
"{{ method.route }}",
{{ method.py_input_message_param }},
{{ method.py_output_message_type.strip('"') }}
{{ method.py_output_message_type.strip('"') }},
timeout=timeout,
deadline=deadline,
metadata=metadata,
)
{% endif %}{# client streaming #}
{% endif %}