Merge pull request #39 from shabbyrobe/reflection-aio

Add support to reflection for aio server

grpc_reflection-stubs/v1alpha/reflection.pyi

@@ -1,14 +1,27 @@
+import typing
+
+import grpc
+from grpc import aio
 from google.protobuf import descriptor_pool
-from grpc import Server, ServicerContext
 from grpc_reflection.v1alpha import reflection_pb2 as _reflection_pb2
 from grpc_reflection.v1alpha._base import BaseReflectionServicer
-from typing import Iterable, List, Optional
 
 SERVICE_NAME: str
 
+AnyServer = typing.Union[grpc.Server, aio.Server]
+AnyServicerContext = typing.Union[grpc.ServicerContext, aio.ServicerContext]
+
 class ReflectionServicer(BaseReflectionServicer):
-    def ServerReflectionInfo(self, request_iterator: Iterable[_reflection_pb2.ServerReflectionRequest], context: ServicerContext) -> None:
+    def ServerReflectionInfo(
+        self,
+        request_iterator: typing.Iterable[_reflection_pb2.ServerReflectionRequest],
+        context: AnyServicerContext,
+    ) -> None:
         ...
 
-def enable_server_reflection(service_names: List[str], server: Server, pool: Optional[descriptor_pool.DescriptorPool] = ...) -> None:
+def enable_server_reflection(
+    service_names: typing.List[str],
+    server: AnyServer,
+    pool: typing.Optional[descriptor_pool.DescriptorPool] = ...,
+) -> None:
     ...

requirements-dev.txt

@@ -1,6 +1,7 @@
 black
 pytest-mypy-plugins==1.10.1
 mypy>=0.902
+types-protobuf
 twine
 wheel
 -e .

typesafety/test_grpc.yml

@@ -206,3 +206,19 @@
     )
     reveal_type(Channel())  # N: Revealed type is "grpc.aio.Channel"
 
+- case: issue_38_reflection_grpc_server
+  main: |
+    import typing
+    import grpc
+    from grpc_reflection.v1alpha.reflection import enable_server_reflection
+    server = typing.cast(grpc.Server, None)
+    enable_server_reflection(["foo"], server, None)
+
+- case: issue_38_reflection_aio_server
+  main: |
+    import typing
+    from grpc import aio
+    from grpc_reflection.v1alpha.reflection import enable_server_reflection
+    server = typing.cast(aio.Server, None)
+    enable_server_reflection(["foo"], server, None)
+