Skip to content

Commit 09f8219

Browse files
nat-nnat-n
committed
Move ServiceStub to a seperate module and add more rpcs to service test
1 parent a757da1 commit 09f8219

File tree

8 files changed

+336
-137
lines changed

8 files changed

+336
-137
lines changed

‎.gitignore

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,4 +9,5 @@ betterproto/tests/output_*
99
dist
1010
**/*.egg-info
1111
output
12-
.idea
12+
.idea
13+
.DS_Store

‎betterproto/__init__.py

Lines changed: 4 additions & 133 deletions
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,9 @@
55
import struct
66
import sys
77
from abc import ABC
8-
from base64 import b64encode, b64decode
8+
from base64 import b64decode, b64encode
99
from datetime import datetime, timedelta, timezone
10+
import stringcase
1011
from typing import (
1112
Any,
1213
AsyncGenerator,
@@ -22,22 +23,12 @@
2223
SupportsBytes,
2324
Tuple,
2425
Type,
25-
TypeVar,
2626
Union,
2727
get_type_hints,
28-
TYPE_CHECKING,
2928
)
30-
31-
32-
import grpclib.const
33-
import stringcase
34-
29+
from ._types import ST, T
3530
from .casing import safe_snake_case
36-
37-
if TYPE_CHECKING:
38-
from grpclib._protocols import IProtoMessage
39-
from grpclib.client import Channel
40-
from grpclib.metadata import Deadline
31+
from .grpc.grpclib_client import ServiceStub
4132

4233
if not (sys.version_info.major == 3 and sys.version_info.minor >= 7):
4334
# Apply backport of datetime.fromisoformat from 3.7
@@ -431,11 +422,6 @@ def parse_fields(value: bytes) -> Generator[ParsedField, None, None]:
431422
)
432423

433424

434-
# Bound type variable to allow methods to return `self` of subclasses
435-
T = TypeVar("T", bound="Message")
436-
ST = TypeVar("ST", bound="IProtoMessage")
437-
438-
439425
class ProtoClassMetadata:
440426
oneof_group_by_field: Dict[str, str]
441427
oneof_field_by_group: Dict[str, Set[dataclasses.Field]]
@@ -1027,118 +1013,3 @@ def _get_wrapper(proto_type: str) -> Type:
10271013
TYPE_STRING: StringValue,
10281014
TYPE_BYTES: BytesValue,
10291015
}[proto_type]
1030-
1031-
1032-
_Value = Union[str, bytes]
1033-
_MetadataLike = Union[Mapping[str, _Value], Collection[Tuple[str, _Value]]]
1034-
1035-
1036-
class ServiceStub(ABC):
1037-
"""
1038-
Base class for async gRPC service stubs.
1039-
"""
1040-
1041-
def __init__(
1042-
self,
1043-
channel: "Channel",
1044-
*,
1045-
timeout: Optional[float] = None,
1046-
deadline: Optional["Deadline"] = None,
1047-
metadata: Optional[_MetadataLike] = None,
1048-
) -> None:
1049-
self.channel = channel
1050-
self.timeout = timeout
1051-
self.deadline = deadline
1052-
self.metadata = metadata
1053-
1054-
def __resolve_request_kwargs(
1055-
self,
1056-
timeout: Optional[float],
1057-
deadline: Optional["Deadline"],
1058-
metadata: Optional[_MetadataLike],
1059-
):
1060-
return {
1061-
"timeout": self.timeout if timeout is None else timeout,
1062-
"deadline": self.deadline if deadline is None else deadline,
1063-
"metadata": self.metadata if metadata is None else metadata,
1064-
}
1065-
1066-
async def _unary_unary(
1067-
self,
1068-
route: str,
1069-
request: "IProtoMessage",
1070-
response_type: Type[T],
1071-
*,
1072-
timeout: Optional[float] = None,
1073-
deadline: Optional["Deadline"] = None,
1074-
metadata: Optional[_MetadataLike] = None,
1075-
) -> T:
1076-
"""Make a unary request and return the response."""
1077-
async with self.channel.request(
1078-
route,
1079-
grpclib.const.Cardinality.UNARY_UNARY,
1080-
type(request),
1081-
response_type,
1082-
**self.__resolve_request_kwargs(timeout, deadline, metadata),
1083-
) as stream:
1084-
await stream.send_message(request, end=True)
1085-
response = await stream.recv_message()
1086-
assert response is not None
1087-
return response
1088-
1089-
async def _unary_stream(
1090-
self,
1091-
route: str,
1092-
request: "IProtoMessage",
1093-
response_type: Type[T],
1094-
*,
1095-
timeout: Optional[float] = None,
1096-
deadline: Optional["Deadline"] = None,
1097-
metadata: Optional[_MetadataLike] = None,
1098-
) -> AsyncGenerator[T, None]:
1099-
"""Make a unary request and return the stream response iterator."""
1100-
async with self.channel.request(
1101-
route,
1102-
grpclib.const.Cardinality.UNARY_STREAM,
1103-
type(request),
1104-
response_type,
1105-
**self.__resolve_request_kwargs(timeout, deadline, metadata),
1106-
) as stream:
1107-
await stream.send_message(request, end=True)
1108-
async for message in stream:
1109-
yield message
1110-
1111-
async def _stream_unary(
1112-
self,
1113-
route: str,
1114-
request_iterator: Iterator["IProtoMessage"],
1115-
request_type: Type[ST],
1116-
response_type: Type[T],
1117-
) -> T:
1118-
"""Make a stream request and return the response."""
1119-
async with self.channel.request(
1120-
route, grpclib.const.Cardinality.STREAM_UNARY, request_type, response_type
1121-
) as stream:
1122-
for message in request_iterator:
1123-
await stream.send_message(message)
1124-
await stream.send_request(end=True)
1125-
response = await stream.recv_message()
1126-
assert response is not None
1127-
return response
1128-
1129-
async def _stream_stream(
1130-
self,
1131-
route: str,
1132-
request_iterator: Iterator["IProtoMessage"],
1133-
request_type: Type[ST],
1134-
response_type: Type[T],
1135-
) -> AsyncGenerator[T, None]:
1136-
"""Make a stream request and return the stream response iterator."""
1137-
async with self.channel.request(
1138-
route, grpclib.const.Cardinality.STREAM_STREAM, request_type, response_type
1139-
) as stream:
1140-
for message in request_iterator:
1141-
await stream.send_message(message)
1142-
await stream.send_request(end=True)
1143-
async for message in stream:
1144-
yield message

‎betterproto/_types.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
from typing import TypeVar
2+
3+
# Bound type variable to allow methods to return `self` of subclasses
4+
T = TypeVar("T", bound="Message")
5+
ST = TypeVar("ST", bound="IProtoMessage")

‎betterproto/grpc/__init__.py

Whitespace-only changes.

‎betterproto/grpc/grpclib_client.py

Lines changed: 135 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,135 @@
1+
from abc import ABC
2+
import grpclib.const
3+
from typing import (
4+
AsyncGenerator,
5+
AsyncIterator,
6+
Collection,
7+
Iterator,
8+
Mapping,
9+
Optional,
10+
Tuple,
11+
TYPE_CHECKING,
12+
Type,
13+
Union,
14+
)
15+
from .._types import ST, T
16+
17+
if TYPE_CHECKING:
18+
from grpclib._protocols import IProtoMessage
19+
from grpclib.client import Channel
20+
from grpclib.metadata import Deadline
21+
22+
23+
_Value = Union[str, bytes]
24+
_MetadataLike = Union[Mapping[str, _Value], Collection[Tuple[str, _Value]]]
25+
26+
27+
class ServiceStub(ABC):
28+
"""
29+
Base class for async gRPC service stubs.
30+
"""
31+
32+
def __init__(
33+
self,
34+
channel: "Channel",
35+
*,
36+
timeout: Optional[float] = None,
37+
deadline: Optional["Deadline"] = None,
38+
metadata: Optional[_MetadataLike] = None,
39+
) -> None:
40+
self.channel = channel
41+
self.timeout = timeout
42+
self.deadline = deadline
43+
self.metadata = metadata
44+
45+
def __resolve_request_kwargs(
46+
self,
47+
timeout: Optional[float],
48+
deadline: Optional["Deadline"],
49+
metadata: Optional[_MetadataLike],
50+
):
51+
return {
52+
"timeout": self.timeout if timeout is None else timeout,
53+
"deadline": self.deadline if deadline is None else deadline,
54+
"metadata": self.metadata if metadata is None else metadata,
55+
}
56+
57+
async def _unary_unary(
58+
self,
59+
route: str,
60+
request: "IProtoMessage",
61+
response_type: Type[T],
62+
*,
63+
timeout: Optional[float] = None,
64+
deadline: Optional["Deadline"] = None,
65+
metadata: Optional[_MetadataLike] = None,
66+
) -> T:
67+
"""Make a unary request and return the response."""
68+
async with self.channel.request(
69+
route,
70+
grpclib.const.Cardinality.UNARY_UNARY,
71+
type(request),
72+
response_type,
73+
**self.__resolve_request_kwargs(timeout, deadline, metadata),
74+
) as stream:
75+
await stream.send_message(request, end=True)
76+
response = await stream.recv_message()
77+
assert response is not None
78+
return response
79+
80+
async def _unary_stream(
81+
self,
82+
route: str,
83+
request: "IProtoMessage",
84+
response_type: Type[T],
85+
*,
86+
timeout: Optional[float] = None,
87+
deadline: Optional["Deadline"] = None,
88+
metadata: Optional[_MetadataLike] = None,
89+
) -> AsyncGenerator[T, None]:
90+
"""Make a unary request and return the stream response iterator."""
91+
async with self.channel.request(
92+
route,
93+
grpclib.const.Cardinality.UNARY_STREAM,
94+
type(request),
95+
response_type,
96+
**self.__resolve_request_kwargs(timeout, deadline, metadata),
97+
) as stream:
98+
await stream.send_message(request, end=True)
99+
async for message in stream:
100+
yield message
101+
102+
async def _stream_unary(
103+
self,
104+
route: str,
105+
request_iterator: Iterator["IProtoMessage"],
106+
request_type: Type[ST],
107+
response_type: Type[T],
108+
) -> T:
109+
"""Make a stream request and return the response."""
110+
async with self.channel.request(
111+
route, grpclib.const.Cardinality.STREAM_UNARY, request_type, response_type
112+
) as stream:
113+
for message in request_iterator:
114+
await stream.send_message(message)
115+
await stream.send_request(end=True)
116+
response = await stream.recv_message()
117+
assert response is not None
118+
return response
119+
120+
async def _stream_stream(
121+
self,
122+
route: str,
123+
request_iterator: Iterator["IProtoMessage"],
124+
request_type: Type[ST],
125+
response_type: Type[T],
126+
) -> AsyncGenerator[T, None]:
127+
"""Make a stream request and return the stream response iterator."""
128+
async with self.channel.request(
129+
route, grpclib.const.Cardinality.STREAM_STREAM, request_type, response_type
130+
) as stream:
131+
for message in request_iterator:
132+
await stream.send_message(message)
133+
await stream.send_request(end=True)
134+
async for message in stream:
135+
yield message

‎betterproto/plugin.py

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -311,7 +311,6 @@ def generate_code(request, response):
311311
}
312312

313313
for j, method in enumerate(service.method):
314-
315314
input_message = None
316315
input_type = get_ref_type(
317316
package, output["imports"], method.input_type

‎betterproto/tests/inputs/service/service.proto

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,13 +3,25 @@ syntax = "proto3";
33
package service;
44

55
message DoThingRequest {
6-
int32 iterations = 1;
6+
string name = 1;
77
}
88

99
message DoThingResponse {
10-
int32 successfulIterations = 1;
10+
repeated string names = 1;
11+
}
12+
13+
message GetThingRequest {
14+
string name = 1;
15+
}
16+
17+
message GetThingResponse {
18+
string name = 1;
19+
int32 version = 2;
1120
}
1221

1322
service Test {
1423
rpc DoThing (DoThingRequest) returns (DoThingResponse);
24+
rpc DoManyThings (stream DoThingRequest) returns (DoThingResponse);
25+
rpc GetThingVersions (GetThingRequest) returns (stream GetThingResponse);
26+
rpc GetDifferentThings (stream GetThingRequest) returns (stream GetThingResponse);
1527
}

0 commit comments

Comments
 (0)