diff --git a/stubs/grpcio/@tests/test_cases/check_rpc_method_handlers.py b/stubs/grpcio/@tests/test_cases/check_rpc_method_handlers.py new file mode 100644 index 000000000000..ae18a55e80e6 --- /dev/null +++ b/stubs/grpcio/@tests/test_cases/check_rpc_method_handlers.py @@ -0,0 +1,288 @@ +from __future__ import annotations + +from collections.abc import AsyncIterator, Callable, Iterator, Mapping +from typing import TYPE_CHECKING, Any, TypeAlias +from typing_extensions import assert_type + +import grpc + + +def sync_unary(request: str, context: grpc.ServicerContext) -> str: + return request + + +async def async_unary(request: str, context: grpc.aio.ServicerContext[str, str]) -> str: + return request + + +async def async_unary_stream(request: str, context: grpc.aio.ServicerContext[str, str]) -> AsyncIterator[str]: + yield request + + +async def async_stream_unary(requests: AsyncIterator[str], context: grpc.aio.ServicerContext[str, str]) -> str: + return await anext(requests) + + +async def async_stream_stream(requests: AsyncIterator[str], context: grpc.aio.ServicerContext[str, str]) -> AsyncIterator[str]: + async for request in requests: + yield request + + +async def async_write(request: str, context: grpc.aio.ServicerContext[str, str]) -> None: + await context.write(request) + + +async def async_stream_write(requests: AsyncIterator[str], context: grpc.aio.ServicerContext[str, str]) -> None: + async for request in requests: + await context.write(request) + + +def sync_unary_stream(request: str, context: grpc.ServicerContext) -> Iterator[str]: + yield request + + +def sync_stream_unary(requests: Iterator[str], context: grpc.ServicerContext) -> str: + return next(requests) + + +def sync_stream(requests: Iterator[str], context: grpc.ServicerContext) -> Iterator[str]: + return requests + + +assert_type(grpc.unary_stream_rpc_method_handler(sync_unary_stream, bytes.decode, str.encode), grpc.RpcMethodHandler[str, str]) +assert_type(grpc.stream_unary_rpc_method_handler(sync_stream_unary, bytes.decode, str.encode), grpc.RpcMethodHandler[str, str]) +assert_type(grpc.unary_unary_rpc_method_handler(sync_unary, bytes.decode, str.encode), grpc.RpcMethodHandler[str, str]) +assert_type(grpc.stream_stream_rpc_method_handler(sync_stream, bytes.decode, str.encode), grpc.RpcMethodHandler[str, str]) +assert_type(grpc.unary_unary_rpc_method_handler(async_unary), grpc.aio._AsyncRpcMethodHandler[str, str]) +assert_type(grpc.unary_stream_rpc_method_handler(async_unary_stream), grpc.aio._AsyncRpcMethodHandler[str, str]) +assert_type(grpc.stream_unary_rpc_method_handler(async_stream_unary), grpc.aio._AsyncRpcMethodHandler[str, str]) +assert_type(grpc.stream_stream_rpc_method_handler(async_stream_stream), grpc.aio._AsyncRpcMethodHandler[str, str]) +assert_type(grpc.unary_stream_rpc_method_handler(async_write), grpc.aio._AsyncRpcMethodHandler[str, str]) +assert_type(grpc.stream_stream_rpc_method_handler(async_stream_write), grpc.aio._AsyncRpcMethodHandler[str, str]) + +sync_handlers = {"Check": grpc.unary_unary_rpc_method_handler(sync_unary, bytes.decode, str.encode)} +grpc.method_handlers_generic_handler("test.Sync", sync_handlers) +async_handlers = {"Check": grpc.unary_unary_rpc_method_handler(async_unary)} +grpc.method_handlers_generic_handler("test.Async", async_handlers) + + +# The synchronous fallback must still accept the callback-based streaming API. +def nonblocking_watch(request: str, context: grpc.ServicerContext, send_response_callback: object) -> None: + pass + + +assert_type(grpc.unary_stream_rpc_method_handler(nonblocking_watch, bytes.decode, str.encode), grpc.RpcMethodHandler[str, str]) + + +# These classes are generic only in the stubs. +if TYPE_CHECKING: + _AsyncGenericHandlerBase: TypeAlias = grpc.GenericRpcHandler[grpc.aio._AsyncRpcMethodHandler[str, str]] + _AsyncServiceHandlerBase: TypeAlias = grpc.ServiceRpcHandler[grpc.aio._AsyncRpcMethodHandler[str, str]] +else: + _AsyncGenericHandlerBase = grpc.GenericRpcHandler + _AsyncServiceHandlerBase = grpc.ServiceRpcHandler + + +# Custom handlers must accept async factory results without widening sync callers. +class AsyncGenericHandler(_AsyncGenericHandlerBase): + def service(self, handler_call_details: grpc.HandlerCallDetails) -> grpc.aio._AsyncRpcMethodHandler[str, str] | None: + if handler_call_details.method == "/test.Async/Check": + return grpc.unary_unary_rpc_method_handler(async_unary) + return None + + +class AsyncServiceHandler(_AsyncServiceHandlerBase): + def service(self, handler_call_details: grpc.HandlerCallDetails) -> grpc.aio._AsyncRpcMethodHandler[str, str]: + return grpc.unary_unary_rpc_method_handler(async_unary) + + def service_name(self) -> str: + return "test.Async" + + +def check_generic_handlers(sync_handler: grpc.GenericRpcHandler, details: grpc.HandlerCallDetails) -> None: + assert_type(sync_handler.service(details), grpc.RpcMethodHandler[Any, Any] | None) + async_handler = grpc.method_handlers_generic_handler("test.Async", async_handlers) + assert_type(async_handler.service(details), grpc.aio._AsyncRpcMethodHandler[str, str] | None) + sync_generic = grpc.method_handlers_generic_handler("test.Sync", sync_handlers) + assert_type(sync_generic.service(details), grpc.RpcMethodHandler[str, str] | None) + server = grpc.aio.server(handlers=[AsyncGenericHandler(), AsyncServiceHandler(), async_handler, sync_handler]) + server.add_generic_rpc_handlers([AsyncGenericHandler(), AsyncServiceHandler(), async_handler, sync_handler]) + + +# An aio server can register and resolve both kinds of handler together. +def check_mixed_registration(server: grpc.aio.Server, details: grpc.HandlerCallDetails) -> None: + mixed: Mapping[str, grpc.RpcMethodHandler[str, str] | grpc.aio._AsyncRpcMethodHandler[str, str]] = { + "Sync": grpc.unary_unary_rpc_method_handler(sync_unary, bytes.decode, str.encode), + "Async": grpc.unary_unary_rpc_method_handler(async_unary), + } + generic = grpc.method_handlers_generic_handler("test.Mixed", mixed) + assert_type(generic.service(details), grpc.RpcMethodHandler[str, str] | grpc.aio._AsyncRpcMethodHandler[str, str] | None) + grpc.aio.server(handlers=[generic]) + server.add_generic_rpc_handlers([generic]) + server.add_registered_method_handlers("test.Mixed", mixed) + + +# Any-returning callbacks must remain usable by synchronous servers. Their +# response type is deliberately not asserted: checkers resolve Any differently. +def dynamic_unary(request: str, context: Any) -> Any: + return request + + +def dynamic_stream(requests: Iterator[str], context: Any) -> Any: + return next(requests) + + +async def partial_async_unary(request: str, context: Any) -> str: + return request + + +async def partial_async_unary_stream(request: str, context: Any) -> AsyncIterator[str]: + yield request + + +async def partial_async_stream_unary(requests: AsyncIterator[str], context: Any) -> str: + return await anext(requests) + + +async def partial_async_stream_stream(requests: AsyncIterator[str], context: Any) -> AsyncIterator[str]: + async for request in requests: + yield request + + +def check_dynamic_callbacks(server: grpc.Server, aio_server: grpc.aio.Server) -> None: + handlers: dict[str, grpc.RpcMethodHandler[str, str]] = { + "UnaryUnary": grpc.unary_unary_rpc_method_handler(dynamic_unary, bytes.decode, str.encode), + "UnaryStream": grpc.unary_stream_rpc_method_handler(dynamic_unary, bytes.decode, str.encode), + "StreamUnary": grpc.stream_unary_rpc_method_handler(dynamic_stream, bytes.decode, str.encode), + "StreamStream": grpc.stream_stream_rpc_method_handler(dynamic_stream, bytes.decode, str.encode), + "KeywordUnaryUnary": grpc.unary_unary_rpc_method_handler(dynamic_unary, response_serializer=str.encode), + "KeywordUnaryStream": grpc.unary_stream_rpc_method_handler(dynamic_unary, response_serializer=str.encode), + "KeywordStreamUnary": grpc.stream_unary_rpc_method_handler(dynamic_stream, response_serializer=str.encode), + "KeywordStreamStream": grpc.stream_stream_rpc_method_handler(dynamic_stream, response_serializer=str.encode), + "RawUnaryUnary": grpc.unary_unary_rpc_method_handler(dynamic_unary), + "RawUnaryStream": grpc.unary_stream_rpc_method_handler(dynamic_unary), + "RawStreamUnary": grpc.stream_unary_rpc_method_handler(dynamic_stream), + "RawStreamStream": grpc.stream_stream_rpc_method_handler(dynamic_stream), + } + server.add_registered_method_handlers("test.Dynamic", handlers) + server.add_generic_rpc_handlers([grpc.method_handlers_generic_handler("test.Dynamic", handlers)]) + aio_server.add_registered_method_handlers("test.Dynamic", handlers) + + +# A broad context annotation must not cause a coroutine to be inferred as a +# synchronous behavior when the serializer supplies a concrete response type. +partial_async_handlers: dict[str, grpc.aio._AsyncRpcMethodHandler[str, str]] = { + "UnaryUnary": grpc.unary_unary_rpc_method_handler(partial_async_unary, bytes.decode, str.encode), + "UnaryStream": grpc.unary_stream_rpc_method_handler(partial_async_unary_stream, bytes.decode, str.encode), + "StreamUnary": grpc.stream_unary_rpc_method_handler(partial_async_stream_unary, bytes.decode, str.encode), + "StreamStream": grpc.stream_stream_rpc_method_handler(partial_async_stream_stream, bytes.decode, str.encode), +} +# Without enough annotations or a serializer, avoid claiming a specific kind. +assert_type(grpc.unary_unary_rpc_method_handler(partial_async_unary), Any) +assert_type(grpc.unary_stream_rpc_method_handler(partial_async_unary_stream), Any) + + +# Optional codecs must not select an async overload for an Any-typed callback. +# Keep factory inference independent of the server registration's expected type. +def check_optional_codecs( + server: grpc.Server, + aio_server: grpc.aio.Server, + callback: Callable[[Any, Any], Any], + serializer: Callable[[str], bytes] | None, + deserializer: Callable[[bytes], str] | None, +) -> None: + uu = grpc.unary_unary_rpc_method_handler(callback, deserializer, serializer) + us = grpc.unary_stream_rpc_method_handler(callback, deserializer, serializer) + su = grpc.stream_unary_rpc_method_handler(callback, deserializer, serializer) + ss = grpc.stream_stream_rpc_method_handler(callback, deserializer, serializer) + assert_type(uu, Any) + assert_type(us, Any) + assert_type(su, Any) + assert_type(ss, Any) + handlers = {"UU": uu, "US": us, "SU": su, "SS": ss} + server.add_registered_method_handlers("test.Optional", handlers) + server.add_generic_rpc_handlers([grpc.method_handlers_generic_handler("test.Optional", handlers)]) + aio_server.add_registered_method_handlers("test.Optional", handlers) + assert_type(grpc.unary_unary_rpc_method_handler(callback, response_serializer=serializer), Any) + assert_type(grpc.unary_stream_rpc_method_handler(callback, response_serializer=serializer), Any) + assert_type(grpc.stream_unary_rpc_method_handler(callback, response_serializer=serializer), Any) + assert_type(grpc.stream_stream_rpc_method_handler(callback, response_serializer=serializer), Any) + + # Explicit aio contexts retain precise types even with optional codecs. + assert_type( + grpc.unary_unary_rpc_method_handler(async_unary, deserializer, serializer), grpc.aio._AsyncRpcMethodHandler[str, str] + ) + assert_type( + grpc.unary_stream_rpc_method_handler(async_unary_stream, deserializer, serializer), + grpc.aio._AsyncRpcMethodHandler[str, str], + ) + assert_type( + grpc.stream_unary_rpc_method_handler(async_stream_unary, deserializer, serializer), + grpc.aio._AsyncRpcMethodHandler[str, str], + ) + assert_type( + grpc.stream_stream_rpc_method_handler(async_stream_stream, deserializer, serializer), + grpc.aio._AsyncRpcMethodHandler[str, str], + ) + assert_type( + grpc.unary_stream_rpc_method_handler(async_write, deserializer, serializer), grpc.aio._AsyncRpcMethodHandler[str, str] + ) + assert_type( + grpc.stream_stream_rpc_method_handler(async_stream_write, deserializer, serializer), + grpc.aio._AsyncRpcMethodHandler[str, str], + ) + + # Broad contexts must not turn coroutine and async-generator results into + # synchronous handlers. Checkers may retain Any for ambiguous signatures. + partial_handlers: dict[str, grpc.aio._AsyncRpcMethodHandler[str, str]] = { + "UU": grpc.unary_unary_rpc_method_handler(partial_async_unary, deserializer, serializer), + "US": grpc.unary_stream_rpc_method_handler(partial_async_unary_stream, deserializer, serializer), + "SU": grpc.stream_unary_rpc_method_handler(partial_async_stream_unary, deserializer, serializer), + "SS": grpc.stream_stream_rpc_method_handler(partial_async_stream_stream, deserializer, serializer), + } + aio_server.add_registered_method_handlers("test.Partial", partial_handlers) + + +# A broad serializer cannot distinguish an awaitable response from a value. +def serialize_object(value: object) -> bytes: + return b"" + + +async def erased_async_unary(request: Any, context: Any) -> str: + return "response" + + +def check_broad_serializers(serializer: Callable[[object], bytes] | None) -> None: + assert_type(grpc.unary_unary_rpc_method_handler(partial_async_unary, bytes.decode, serialize_object), Any) + assert_type(grpc.unary_unary_rpc_method_handler(partial_async_unary, response_serializer=serialize_object), Any) + assert_type(grpc.stream_unary_rpc_method_handler(erased_async_unary, bytes.decode, serialize_object), Any) + assert_type(grpc.stream_unary_rpc_method_handler(erased_async_unary, response_serializer=serialize_object), Any) + assert_type(grpc.unary_unary_rpc_method_handler(partial_async_unary, bytes.decode, serializer), Any) + assert_type(grpc.stream_unary_rpc_method_handler(erased_async_unary, bytes.decode, serializer), Any) + assert_type( + grpc.unary_unary_rpc_method_handler(async_unary, bytes.decode, serialize_object), + grpc.aio._AsyncRpcMethodHandler[str, str], + ) + assert_type( + grpc.stream_unary_rpc_method_handler(async_stream_unary, bytes.decode, serialize_object), + grpc.aio._AsyncRpcMethodHandler[str, str], + ) + + +# Synchronous-only contexts must retain their handler interface without codecs +# and when a codec is optional, rather than falling through to Any. +def check_sync_inference(serializer: Callable[[str], bytes] | None) -> None: + assert_type(grpc.unary_unary_rpc_method_handler(sync_unary), grpc.RpcMethodHandler[str, str]) + assert_type(grpc.unary_stream_rpc_method_handler(sync_unary_stream), grpc.RpcMethodHandler[str, str]) + assert_type(grpc.stream_unary_rpc_method_handler(sync_stream_unary), grpc.RpcMethodHandler[str, str]) + assert_type(grpc.stream_stream_rpc_method_handler(sync_stream), grpc.RpcMethodHandler[str, str]) + assert_type(grpc.unary_unary_rpc_method_handler(sync_unary, response_serializer=serializer), grpc.RpcMethodHandler[str, str]) + assert_type( + grpc.unary_stream_rpc_method_handler(sync_unary_stream, response_serializer=serializer), grpc.RpcMethodHandler[str, str] + ) + assert_type( + grpc.stream_unary_rpc_method_handler(sync_stream_unary, response_serializer=serializer), grpc.RpcMethodHandler[str, str] + ) + assert_type( + grpc.stream_stream_rpc_method_handler(sync_stream, response_serializer=serializer), grpc.RpcMethodHandler[str, str] + ) diff --git a/stubs/grpcio/@tests/test_cases/check_server_interceptor.py b/stubs/grpcio/@tests/test_cases/check_server_interceptor.py index febad80e4bb8..2a365e6aafde 100644 --- a/stubs/grpcio/@tests/test_cases/check_server_interceptor.py +++ b/stubs/grpcio/@tests/test_cases/check_server_interceptor.py @@ -1,12 +1,17 @@ from __future__ import annotations +import inspect from collections.abc import Awaitable, Callable from concurrent.futures.thread import ThreadPoolExecutor -from typing import TypeVar +from typing import TYPE_CHECKING, TypeVar +from typing_extensions import TypeGuard import grpc import grpc.aio +if TYPE_CHECKING: + from grpc.aio import _AsyncRpcMethodHandler + RequestT = TypeVar("RequestT") ResponseT = TypeVar("ResponseT") @@ -26,10 +31,39 @@ def intercept_service( class NoopAioInterceptor(grpc.aio.ServerInterceptor): async def intercept_service( self, - continuation: Callable[[grpc.HandlerCallDetails], Awaitable[grpc.RpcMethodHandler[RequestT, ResponseT] | None]], + continuation: Callable[ + [grpc.HandlerCallDetails], + Awaitable[grpc.RpcMethodHandler[RequestT, ResponseT] | _AsyncRpcMethodHandler[RequestT, ResponseT] | None], + ], handler_call_details: grpc.HandlerCallDetails, - ) -> grpc.RpcMethodHandler[RequestT, ResponseT] | None: + ) -> grpc.RpcMethodHandler[RequestT, ResponseT] | _AsyncRpcMethodHandler[RequestT, ResponseT] | None: return await continuation(handler_call_details) grpc.aio.server(interceptors=[NoopAioInterceptor()]) + + +def is_async_unary_handler( + handler: grpc.RpcMethodHandler[RequestT, ResponseT] | _AsyncRpcMethodHandler[RequestT, ResponseT], +) -> TypeGuard[_AsyncRpcMethodHandler[RequestT, ResponseT]]: + return inspect.iscoroutinefunction(handler.unary_unary) + + +class WrappingAioInterceptor(grpc.aio.ServerInterceptor): + async def intercept_service( + self, + continuation: Callable[ + [grpc.HandlerCallDetails], + Awaitable[grpc.RpcMethodHandler[RequestT, ResponseT] | _AsyncRpcMethodHandler[RequestT, ResponseT] | None], + ], + handler_call_details: grpc.HandlerCallDetails, + ) -> grpc.RpcMethodHandler[RequestT, ResponseT] | _AsyncRpcMethodHandler[RequestT, ResponseT] | None: + handler = await continuation(handler_call_details) + if handler is None or not is_async_unary_handler(handler) or handler.unary_unary is None: + return handler + inner = handler.unary_unary + + async def behavior(request: RequestT, context: grpc.aio.ServicerContext[RequestT, ResponseT]) -> ResponseT: + return await inner(request, context) + + return grpc.unary_unary_rpc_method_handler(behavior, handler.request_deserializer, handler.response_serializer) diff --git a/stubs/grpcio/grpc/__init__.pyi b/stubs/grpcio/grpc/__init__.pyi index 9e7160589f9c..623a05436d9d 100644 --- a/stubs/grpcio/grpc/__init__.pyi +++ b/stubs/grpcio/grpc/__init__.pyi @@ -1,11 +1,11 @@ import abc import enum import threading -from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence +from collections.abc import AsyncIterator, Awaitable, Callable, Iterable, Iterator, Mapping, Sequence from concurrent import futures from types import ModuleType, TracebackType -from typing import Any, Generic, Protocol, TypeAlias, TypeVar, runtime_checkable, type_check_only -from typing_extensions import Never, Self +from typing import Any, Generic, Protocol, TypeAlias, overload, runtime_checkable, type_check_only +from typing_extensions import Never, Self, TypeVar from . import aio as aio @@ -38,6 +38,12 @@ _Metadata: TypeAlias = tuple[tuple[str, str | bytes], ...] _TRequest = TypeVar("_TRequest") _TResponse = TypeVar("_TResponse") +_TMethodHandler_co = TypeVar( + "_TMethodHandler_co", + bound=RpcMethodHandler[Any, Any] | aio._AsyncRpcMethodHandler[Any, Any], + default=RpcMethodHandler[Any, Any], + covariant=True, +) _Serializer: TypeAlias = Callable[[_T], bytes] _Deserializer: TypeAlias = Callable[[bytes], _T] @@ -148,29 +154,118 @@ def xds_server_credentials(fallback_credentials: ServerCredentials) -> ServerCre class _Behaviour(Protocol): def __call__(self, *args: Any, **kwargs: Any) -> Any: ... +# Callbacks accepting both contexts may be synchronous or asynchronous. +# Keep their handler type gradual; serializers can accept arbitrary objects +# and cannot reliably distinguish values from awaitables. +@overload +def unary_unary_rpc_method_handler( + behavior: Callable[[_TRequest, ServicerContext | aio.ServicerContext[_TRequest, _TResponse]], Any], + request_deserializer: _Deserializer[_TRequest] | None = None, + response_serializer: _Serializer[_TResponse] | None = None, +) -> Any: ... +@overload +def unary_unary_rpc_method_handler( + behavior: Callable[[_TRequest, ServicerContext], _TResponse], + request_deserializer: _Deserializer[_TRequest] | None = None, + response_serializer: _Serializer[_TResponse] | None = None, +) -> RpcMethodHandler[_TRequest, _TResponse]: ... +@overload +def unary_unary_rpc_method_handler( + behavior: Callable[[_TRequest, aio.ServicerContext[_TRequest, _TResponse]], Awaitable[_TResponse]], + request_deserializer: _Deserializer[_TRequest] | None = None, + response_serializer: _Serializer[_TResponse] | None = None, +) -> aio._AsyncRpcMethodHandler[_TRequest, _TResponse]: ... +@overload def unary_unary_rpc_method_handler( behavior: _Behaviour, request_deserializer: _Deserializer[_TRequest] | None = None, response_serializer: _Serializer[_TResponse] | None = None, ) -> RpcMethodHandler[_TRequest, _TResponse]: ... + +@overload +def unary_stream_rpc_method_handler( + behavior: Callable[[_TRequest, ServicerContext | aio.ServicerContext[_TRequest, _TResponse]], Any], + request_deserializer: _Deserializer[_TRequest] | None = None, + response_serializer: _Serializer[_TResponse] | None = None, +) -> Any: ... +@overload +def unary_stream_rpc_method_handler( + behavior: Callable[[_TRequest, ServicerContext], Iterator[_TResponse]], + request_deserializer: _Deserializer[_TRequest] | None = None, + response_serializer: _Serializer[_TResponse] | None = None, +) -> RpcMethodHandler[_TRequest, _TResponse]: ... +@overload +def unary_stream_rpc_method_handler( + behavior: Callable[[_TRequest, aio.ServicerContext[_TRequest, _TResponse]], AsyncIterator[_TResponse] | Awaitable[None]], + request_deserializer: _Deserializer[_TRequest] | None = None, + response_serializer: _Serializer[_TResponse] | None = None, +) -> aio._AsyncRpcMethodHandler[_TRequest, _TResponse]: ... +@overload def unary_stream_rpc_method_handler( behavior: _Behaviour, request_deserializer: _Deserializer[_TRequest] | None = None, response_serializer: _Serializer[_TResponse] | None = None, ) -> RpcMethodHandler[_TRequest, _TResponse]: ... + +@overload +def stream_unary_rpc_method_handler( + behavior: Callable[ + [Iterator[_TRequest] | AsyncIterator[_TRequest], ServicerContext | aio.ServicerContext[_TRequest, _TResponse]], Any + ], + request_deserializer: _Deserializer[_TRequest] | None = None, + response_serializer: _Serializer[_TResponse] | None = None, +) -> Any: ... +@overload +def stream_unary_rpc_method_handler( + behavior: Callable[[Iterator[_TRequest], ServicerContext], _TResponse], + request_deserializer: _Deserializer[_TRequest] | None = None, + response_serializer: _Serializer[_TResponse] | None = None, +) -> RpcMethodHandler[_TRequest, _TResponse]: ... +@overload +def stream_unary_rpc_method_handler( + behavior: Callable[[AsyncIterator[_TRequest], aio.ServicerContext[_TRequest, _TResponse]], Awaitable[_TResponse]], + request_deserializer: _Deserializer[_TRequest] | None = None, + response_serializer: _Serializer[_TResponse] | None = None, +) -> aio._AsyncRpcMethodHandler[_TRequest, _TResponse]: ... +@overload def stream_unary_rpc_method_handler( behavior: _Behaviour, request_deserializer: _Deserializer[_TRequest] | None = None, response_serializer: _Serializer[_TResponse] | None = None, ) -> RpcMethodHandler[_TRequest, _TResponse]: ... + +@overload +def stream_stream_rpc_method_handler( + behavior: Callable[ + [Iterator[_TRequest] | AsyncIterator[_TRequest], ServicerContext | aio.ServicerContext[_TRequest, _TResponse]], Any + ], + request_deserializer: _Deserializer[_TRequest] | None = None, + response_serializer: _Serializer[_TResponse] | None = None, +) -> Any: ... +@overload +def stream_stream_rpc_method_handler( + behavior: Callable[[Iterator[_TRequest], ServicerContext], Iterator[_TResponse]], + request_deserializer: _Deserializer[_TRequest] | None = None, + response_serializer: _Serializer[_TResponse] | None = None, +) -> RpcMethodHandler[_TRequest, _TResponse]: ... +@overload +def stream_stream_rpc_method_handler( + behavior: Callable[ + [AsyncIterator[_TRequest], aio.ServicerContext[_TRequest, _TResponse]], AsyncIterator[_TResponse] | Awaitable[None] + ], + request_deserializer: _Deserializer[_TRequest] | None = None, + response_serializer: _Serializer[_TResponse] | None = None, +) -> aio._AsyncRpcMethodHandler[_TRequest, _TResponse]: ... +@overload def stream_stream_rpc_method_handler( behavior: _Behaviour, request_deserializer: _Deserializer[_TRequest] | None = None, response_serializer: _Serializer[_TResponse] | None = None, ) -> RpcMethodHandler[_TRequest, _TResponse]: ... + def method_handlers_generic_handler( - service: str, method_handlers: dict[str, RpcMethodHandler[Any, Any]] -) -> GenericRpcHandler: ... + service: str, method_handlers: Mapping[str, _TMethodHandler_co] +) -> GenericRpcHandler[_TMethodHandler_co]: ... # Channel Ready Future: @@ -491,12 +586,12 @@ class HandlerCallDetails(Protocol): method: str invocation_metadata: _Metadata -class GenericRpcHandler(abc.ABC): +class GenericRpcHandler(abc.ABC, Generic[_TMethodHandler_co]): # The return type depends on the handler call details. @abc.abstractmethod - def service(self, handler_call_details: HandlerCallDetails) -> RpcMethodHandler[Any, Any] | None: ... + def service(self, handler_call_details: HandlerCallDetails) -> _TMethodHandler_co | None: ... -class ServiceRpcHandler(GenericRpcHandler, metaclass=abc.ABCMeta): +class ServiceRpcHandler(GenericRpcHandler[_TMethodHandler_co], metaclass=abc.ABCMeta): @abc.abstractmethod def service_name(self) -> str: ... diff --git a/stubs/grpcio/grpc/aio/__init__.pyi b/stubs/grpcio/grpc/aio/__init__.pyi index 06819e6aa75a..30e824ace605 100644 --- a/stubs/grpcio/grpc/aio/__init__.pyi +++ b/stubs/grpcio/grpc/aio/__init__.pyi @@ -129,7 +129,7 @@ def secure_channel( def server( migration_thread_pool: futures.Executor | None = None, - handlers: Sequence[GenericRpcHandler] | None = None, + handlers: Sequence[GenericRpcHandler[RpcMethodHandler[Any, Any] | _AsyncRpcMethodHandler[Any, Any]]] | None = None, interceptors: Sequence[ServerInterceptor] | None = None, options: _Options | None = None, maximum_concurrent_rpcs: int | None = None, @@ -189,9 +189,11 @@ class Channel(abc.ABC): class Server(metaclass=abc.ABCMeta): @abc.abstractmethod - def add_generic_rpc_handlers(self, generic_rpc_handlers: Iterable[GenericRpcHandler]) -> None: ... + def add_generic_rpc_handlers( + self, generic_rpc_handlers: Iterable[GenericRpcHandler[RpcMethodHandler[Any, Any] | _AsyncRpcMethodHandler[Any, Any]]] + ) -> None: ... def add_registered_method_handlers( - self, service_name: str, method_handlers: Mapping[str, RpcMethodHandler[Any, Any]] + self, service_name: str, method_handlers: Mapping[str, RpcMethodHandler[Any, Any] | _AsyncRpcMethodHandler[Any, Any]] ) -> None: ... # Returns an integer port on which server will accept RPC requests. @@ -443,7 +445,27 @@ class StreamStreamClientInterceptor(ClientInterceptor, metaclass=abc.ABCMeta): # Server-Side Interceptor: +# Async version of grpc.RpcMethodHandler. As opposed to the former, this type +# does not exist at runtime. +@type_check_only +class _AsyncRpcMethodHandler(Generic[_TRequest, _TResponse]): + request_streaming: bool + response_streaming: bool + request_deserializer: _Deserializer[_TRequest] | None + response_serializer: _Serializer[_TResponse] | None + unary_unary: Callable[[_TRequest, ServicerContext[_TRequest, _TResponse]], Awaitable[_TResponse]] | None + unary_stream: ( + Callable[[_TRequest, ServicerContext[_TRequest, _TResponse]], AsyncIterator[_TResponse] | Awaitable[None]] | None + ) + stream_unary: Callable[[AsyncIterator[_TRequest], ServicerContext[_TRequest, _TResponse]], Awaitable[_TResponse]] | None + stream_stream: ( + Callable[[AsyncIterator[_TRequest], ServicerContext[_TRequest, _TResponse]], AsyncIterator[_TResponse] | Awaitable[None]] + | None + ) + class ServerInterceptor(metaclass=abc.ABCMeta): + # aio servers also run synchronous handlers in a thread pool, so continuation + # can return either kind of handler without adapting its behavior. # This method (not the class) is generic over _TRequest and _TResponse # and the types must satisfy the no-op implementation of # `return await continuation(handler_call_details)`. @@ -453,9 +475,12 @@ class ServerInterceptor(metaclass=abc.ABCMeta): @abc.abstractmethod async def intercept_service( self, - continuation: Callable[[HandlerCallDetails], Awaitable[RpcMethodHandler[_TRequest, _TResponse] | None]], + continuation: Callable[ + [HandlerCallDetails], + Awaitable[RpcMethodHandler[_TRequest, _TResponse] | _AsyncRpcMethodHandler[_TRequest, _TResponse] | None], + ], handler_call_details: HandlerCallDetails, - ) -> RpcMethodHandler[_TRequest, _TResponse] | None: ... + ) -> RpcMethodHandler[_TRequest, _TResponse] | _AsyncRpcMethodHandler[_TRequest, _TResponse] | None: ... # Multi-Callable Interfaces: