Skip to content

Commit c5b10d7

Browse files
RobinMcCorkellRobinMcCorkell
committed
grpcio: improve client interceptors
1 parent 725a900 commit c5b10d7

File tree

3 files changed

+154
-42
lines changed

3 files changed

+154
-42
lines changed

‎stubs/grpcio/@tests/test_cases/check_client_interceptor.py‎

Lines changed: 117 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,117 @@
1+
from __future__ import annotations
2+
3+
from collections.abc import Callable
4+
from typing import AsyncIterable, AsyncIterator, Awaitable, Iterable, Iterator, TypeVar
5+
6+
import grpc
7+
import grpc.aio
8+
9+
RequestT = TypeVar("RequestT")
10+
ResponseT = TypeVar("ResponseT")
11+
12+
13+
class NoopUnaryUnaryInterceptor(grpc.UnaryUnaryClientInterceptor):
14+
def intercept_unary_unary(
15+
self,
16+
continuation: Callable[[grpc.ClientCallDetails, RequestT], grpc._CallFuture[ResponseT]],
17+
client_call_details: grpc.ClientCallDetails,
18+
request: RequestT,
19+
) -> grpc._CallFuture[ResponseT]:
20+
return continuation(client_call_details, request)
21+
22+
23+
class NoopUnaryStreamInterceptor(grpc.UnaryStreamClientInterceptor):
24+
def intercept_unary_stream(
25+
self,
26+
continuation: Callable[[grpc.ClientCallDetails, RequestT], grpc._CallIterator[ResponseT]],
27+
client_call_details: grpc.ClientCallDetails,
28+
request: RequestT,
29+
) -> grpc._CallIterator[ResponseT]:
30+
return continuation(client_call_details, request)
31+
32+
33+
class NoopStreamUnaryInterceptor(grpc.StreamUnaryClientInterceptor):
34+
def intercept_stream_unary(
35+
self,
36+
continuation: Callable[[grpc.ClientCallDetails, Iterator[RequestT]], grpc._CallFuture[ResponseT]],
37+
client_call_details: grpc.ClientCallDetails,
38+
request_iterator: Iterator[RequestT],
39+
) -> grpc._CallFuture[ResponseT]:
40+
return continuation(client_call_details, request_iterator)
41+
42+
43+
class NoopStreamStreamInterceptor(grpc.StreamStreamClientInterceptor):
44+
def intercept_stream_stream(
45+
self,
46+
continuation: Callable[[grpc.ClientCallDetails, Iterator[RequestT]], grpc._CallIterator[ResponseT]],
47+
client_call_details: grpc.ClientCallDetails,
48+
request_iterator: Iterator[RequestT],
49+
) -> grpc._CallIterator[ResponseT]:
50+
return continuation(client_call_details, request_iterator)
51+
52+
53+
channel = grpc.insecure_channel("target")
54+
channel = grpc.intercept_channel(
55+
channel,
56+
NoopUnaryUnaryInterceptor(),
57+
NoopUnaryStreamInterceptor(),
58+
NoopStreamUnaryInterceptor(),
59+
NoopStreamStreamInterceptor(),
60+
)
61+
62+
63+
class NoopAioUnaryUnaryInterceptor(grpc.aio.UnaryUnaryClientInterceptor):
64+
async def intercept_unary_unary(
65+
self,
66+
continuation: Callable[[grpc.aio.ClientCallDetails, RequestT], Awaitable[grpc.aio.UnaryUnaryCall[RequestT, ResponseT]]],
67+
client_call_details: grpc.aio.ClientCallDetails,
68+
request: RequestT,
69+
) -> ResponseT | grpc.aio.UnaryUnaryCall[RequestT, ResponseT]:
70+
return await continuation(client_call_details, request)
71+
72+
73+
class NoopAioUnaryStreamInterceptor(grpc.aio.UnaryStreamClientInterceptor):
74+
async def intercept_unary_stream(
75+
self,
76+
continuation: Callable[[grpc.aio.ClientCallDetails, RequestT], Awaitable[grpc.aio.UnaryStreamCall[RequestT, ResponseT]]],
77+
client_call_details: grpc.aio.ClientCallDetails,
78+
request: RequestT,
79+
) -> AsyncIterator[ResponseT] | grpc.aio.UnaryStreamCall[RequestT, ResponseT]:
80+
return await continuation(client_call_details, request)
81+
82+
83+
class NoopAioStreamUnaryInterceptor(grpc.aio.StreamUnaryClientInterceptor):
84+
async def intercept_stream_unary(
85+
self,
86+
continuation: Callable[
87+
[grpc.aio.ClientCallDetails, AsyncIterable[RequestT] | Iterable[RequestT]],
88+
Awaitable[grpc.aio.StreamUnaryCall[RequestT, ResponseT]],
89+
],
90+
client_call_details: grpc.aio.ClientCallDetails,
91+
request_iterator: AsyncIterable[RequestT] | Iterable[RequestT],
92+
) -> ResponseT | grpc.aio.StreamUnaryCall[RequestT, ResponseT]:
93+
return await continuation(client_call_details, request_iterator)
94+
95+
96+
class NoopAioStreamStreamInterceptor(grpc.aio.StreamStreamClientInterceptor):
97+
async def intercept_stream_stream(
98+
self,
99+
continuation: Callable[
100+
[grpc.aio.ClientCallDetails, AsyncIterable[RequestT] | Iterable[RequestT]],
101+
Awaitable[grpc.aio.StreamStreamCall[RequestT, ResponseT]],
102+
],
103+
client_call_details: grpc.aio.ClientCallDetails,
104+
request_iterator: AsyncIterable[RequestT] | Iterable[RequestT],
105+
) -> AsyncIterator[ResponseT] | grpc.aio.StreamStreamCall[RequestT, ResponseT]:
106+
return await continuation(client_call_details, request_iterator)
107+
108+
109+
grpc.aio.insecure_channel(
110+
"target",
111+
interceptors=[
112+
NoopAioUnaryUnaryInterceptor(),
113+
NoopAioUnaryStreamInterceptor(),
114+
NoopAioStreamUnaryInterceptor(),
115+
NoopAioStreamStreamInterceptor(),
116+
],
117+
)

‎stubs/grpcio/grpc/__init__.pyi‎

Lines changed: 12 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -76,13 +76,10 @@ def secure_channel(
7676
) -> Channel: ...
7777

7878
_Interceptor: TypeAlias = (
79-
UnaryUnaryClientInterceptor[_TRequest, _TResponse]
80-
| UnaryStreamClientInterceptor[_TRequest, _TResponse]
81-
| StreamUnaryClientInterceptor[_TRequest, _TResponse]
82-
| StreamStreamClientInterceptor[_TRequest, _TResponse]
79+
UnaryUnaryClientInterceptor | UnaryStreamClientInterceptor | StreamUnaryClientInterceptor | StreamStreamClientInterceptor
8380
)
8481

85-
def intercept_channel(channel: Channel, *interceptors: _Interceptor[_TRequest, _TResponse]) -> Channel: ...
82+
def intercept_channel(channel: Channel, *interceptors: _Interceptor) -> Channel: ...
8683

8784
# Create Client Credentials:
8885

@@ -378,25 +375,11 @@ class ClientCallDetails(abc.ABC):
378375
@type_check_only
379376
class _CallFuture(Call, Future[_TResponse], metaclass=abc.ABCMeta): ...
380377

381-
class UnaryUnaryClientInterceptor(abc.ABC, Generic[_TRequest, _TResponse]):
378+
class UnaryUnaryClientInterceptor(abc.ABC):
379+
# This method (not the class) is generic over _TRequest and _TResponse.
382380
@abc.abstractmethod
383381
def intercept_unary_unary(
384382
self,
385-
# FIXME: decode these cryptic runes to confirm the typing mystery of
386-
# this callable's signature that was left for us by past civilisations:
387-
#
388-
# continuation - A function that proceeds with the invocation by
389-
# executing the next interceptor in chain or invoking the actual RPC
390-
# on the underlying Channel. It is the interceptor's responsibility
391-
# to call it if it decides to move the RPC forward. The interceptor
392-
# can use response_future = continuation(client_call_details,
393-
# request) to continue with the RPC. continuation returns an object
394-
# that is both a Call for the RPC and a Future. In the event of RPC
395-
# completion, the return Call-Future's result value will be the
396-
# response message of the RPC. Should the event terminate with non-OK
397-
# status, the returned Call-Future's exception value will be an
398-
# RpcError.
399-
#
400383
continuation: Callable[[ClientCallDetails, _TRequest], _CallFuture[_TResponse]],
401384
client_call_details: ClientCallDetails,
402385
request: _TRequest,
@@ -407,7 +390,8 @@ class _CallIterator(Call, Generic[_TResponse], metaclass=abc.ABCMeta):
407390
def __iter__(self) -> Iterator[_TResponse]: ...
408391
def __next__(self) -> _TResponse: ...
409392

410-
class UnaryStreamClientInterceptor(abc.ABC, Generic[_TRequest, _TResponse]):
393+
class UnaryStreamClientInterceptor(abc.ABC):
394+
# This method (not the class) is generic over _TRequest and _TResponse.
411395
@abc.abstractmethod
412396
def intercept_unary_stream(
413397
self,
@@ -416,20 +400,22 @@ class UnaryStreamClientInterceptor(abc.ABC, Generic[_TRequest, _TResponse]):
416400
request: _TRequest,
417401
) -> _CallIterator[_TResponse]: ...
418402

419-
class StreamUnaryClientInterceptor(abc.ABC, Generic[_TRequest, _TResponse]):
403+
class StreamUnaryClientInterceptor(abc.ABC):
404+
# This method (not the class) is generic over _TRequest and _TResponse.
420405
@abc.abstractmethod
421406
def intercept_stream_unary(
422407
self,
423-
continuation: Callable[[ClientCallDetails, _TRequest], _CallFuture[_TResponse]],
408+
continuation: Callable[[ClientCallDetails, Iterator[_TRequest]], _CallFuture[_TResponse]],
424409
client_call_details: ClientCallDetails,
425410
request_iterator: Iterator[_TRequest],
426411
) -> _CallFuture[_TResponse]: ...
427412

428-
class StreamStreamClientInterceptor(abc.ABC, Generic[_TRequest, _TResponse]):
413+
class StreamStreamClientInterceptor(abc.ABC):
414+
# This method (not the class) is generic over _TRequest and _TResponse.
429415
@abc.abstractmethod
430416
def intercept_stream_stream(
431417
self,
432-
continuation: Callable[[ClientCallDetails, _TRequest], _CallIterator[_TResponse]],
418+
continuation: Callable[[ClientCallDetails, Iterator[_TRequest]], _CallIterator[_TResponse]],
433419
client_call_details: ClientCallDetails,
434420
request_iterator: Iterator[_TRequest],
435421
) -> _CallIterator[_TResponse]: ...

‎stubs/grpcio/grpc/aio/__init__.pyi‎

Lines changed: 25 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -45,8 +45,6 @@ class AioRpcError(RpcError):
4545

4646
# Create Client:
4747

48-
class ClientInterceptor(metaclass=abc.ABCMeta): ...
49-
5048
def insecure_channel(
5149
target: str,
5250
options: _Options | None = None,
@@ -288,7 +286,7 @@ class InterceptedUnaryUnaryCall(_InterceptedCall[_TRequest, _TResponse], metacla
288286
def __await__(self) -> Generator[Incomplete, None, _TResponse]: ...
289287
def __init__(
290288
self,
291-
interceptors: Sequence[UnaryUnaryClientInterceptor[_TRequest, _TResponse]],
289+
interceptors: Sequence[UnaryUnaryClientInterceptor],
292290
request: _TRequest,
293291
timeout: float | None,
294292
metadata: Metadata,
@@ -304,7 +302,7 @@ class InterceptedUnaryUnaryCall(_InterceptedCall[_TRequest, _TResponse], metacla
304302
# pylint: disable=too-many-arguments
305303
async def _invoke(
306304
self,
307-
interceptors: Sequence[UnaryUnaryClientInterceptor[_TRequest, _TResponse]],
305+
interceptors: Sequence[UnaryUnaryClientInterceptor],
308306
method: bytes,
309307
timeout: float | None,
310308
metadata: Metadata | None,
@@ -316,42 +314,53 @@ class InterceptedUnaryUnaryCall(_InterceptedCall[_TRequest, _TResponse], metacla
316314
) -> UnaryUnaryCall[_TRequest, _TResponse]: ...
317315
def time_remaining(self) -> float | None: ...
318316

319-
class UnaryUnaryClientInterceptor(Generic[_TRequest, _TResponse], metaclass=abc.ABCMeta):
317+
class ClientInterceptor(metaclass=abc.ABCMeta): ...
318+
319+
class UnaryUnaryClientInterceptor(ClientInterceptor, metaclass=abc.ABCMeta):
320+
# This method (not the class) is generic over _TRequest and _TResponse.
320321
@abc.abstractmethod
321322
async def intercept_unary_unary(
322323
self,
323324
# XXX: See equivalent function in grpc types for notes about continuation:
324-
continuation: Callable[[ClientCallDetails, _TRequest], UnaryUnaryCall[_TRequest, _TResponse]],
325+
continuation: Callable[[ClientCallDetails, _TRequest], Awaitable[UnaryUnaryCall[_TRequest, _TResponse]]],
325326
client_call_details: ClientCallDetails,
326327
request: _TRequest,
327-
) -> _TResponse: ...
328+
) -> _TResponse | UnaryUnaryCall[_TRequest, _TResponse]: ...
328329

329-
class UnaryStreamClientInterceptor(Generic[_TRequest, _TResponse], metaclass=abc.ABCMeta):
330+
class UnaryStreamClientInterceptor(ClientInterceptor, metaclass=abc.ABCMeta):
331+
# This method (not the class) is generic over _TRequest and _TResponse.
330332
@abc.abstractmethod
331333
async def intercept_unary_stream(
332334
self,
333-
continuation: Callable[[ClientCallDetails, _TRequest], UnaryStreamCall[_TRequest, _TResponse]],
335+
continuation: Callable[[ClientCallDetails, _TRequest], Awaitable[UnaryStreamCall[_TRequest, _TResponse]]],
334336
client_call_details: ClientCallDetails,
335337
request: _TRequest,
336-
) -> AsyncIterable[_TResponse] | UnaryStreamCall[_TRequest, _TResponse]: ...
338+
) -> AsyncIterator[_TResponse] | UnaryStreamCall[_TRequest, _TResponse]: ...
337339

338-
class StreamUnaryClientInterceptor(Generic[_TRequest, _TResponse], metaclass=abc.ABCMeta):
340+
class StreamUnaryClientInterceptor(ClientInterceptor, metaclass=abc.ABCMeta):
341+
# This method (not the class) is generic over _TRequest and _TResponse.
339342
@abc.abstractmethod
340343
async def intercept_stream_unary(
341344
self,
342-
continuation: Callable[[ClientCallDetails, _TRequest], StreamUnaryCall[_TRequest, _TResponse]],
345+
continuation: Callable[
346+
[ClientCallDetails, AsyncIterable[_TRequest] | Iterable[_TRequest]], Awaitable[StreamUnaryCall[_TRequest, _TResponse]]
347+
],
343348
client_call_details: ClientCallDetails,
344349
request_iterator: AsyncIterable[_TRequest] | Iterable[_TRequest],
345-
) -> AsyncIterable[_TResponse] | UnaryStreamCall[_TRequest, _TResponse]: ...
350+
) -> _TResponse | StreamUnaryCall[_TRequest, _TResponse]: ...
346351

347-
class StreamStreamClientInterceptor(Generic[_TRequest, _TResponse], metaclass=abc.ABCMeta):
352+
class StreamStreamClientInterceptor(ClientInterceptor, metaclass=abc.ABCMeta):
353+
# This method (not the class) is generic over _TRequest and _TResponse.
348354
@abc.abstractmethod
349355
async def intercept_stream_stream(
350356
self,
351-
continuation: Callable[[ClientCallDetails, _TRequest], StreamStreamCall[_TRequest, _TResponse]],
357+
continuation: Callable[
358+
[ClientCallDetails, AsyncIterable[_TRequest] | Iterable[_TRequest]],
359+
Awaitable[StreamStreamCall[_TRequest, _TResponse]],
360+
],
352361
client_call_details: ClientCallDetails,
353362
request_iterator: AsyncIterable[_TRequest] | Iterable[_TRequest],
354-
) -> AsyncIterable[_TResponse] | StreamStreamCall[_TRequest, _TResponse]: ...
363+
) -> AsyncIterator[_TResponse] | StreamStreamCall[_TRequest, _TResponse]: ...
355364

356365
# Server-Side Interceptor:
357366

0 commit comments

Comments
 (0)