Fix UnaryStreamMultiCallable and StreamStreamMultiCallable Future/Call

shabbyrobe · May 25, 2023 · e554254 · e554254
1 parent b2d5e65
commit e554254
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 2 deletions.
diff --git a/grpc-stubs/__init__.pyi b/grpc-stubs/__init__.pyi
@@ -669,7 +669,7 @@ class UnaryUnaryMultiCallable(typing.Generic[TRequest, TResponse]):
         ...
 
 
-class UnaryStreamMultiCallable(typing.Generic[TRequest, TResponse]):
+class UnaryStreamMultiCallable(CallFuture[TResponse], typing.Generic[TRequest, TResponse]):
     def __call__(
         self,
         request: TRequest,
@@ -733,7 +733,7 @@ class StreamUnaryMultiCallable(typing.Generic[TRequest, TResponse]):
         ...
 
 
-class StreamStreamMultiCallable(typing.Generic[TRequest, TResponse]):
+class StreamStreamMultiCallable(CallFuture[TResponse], typing.Generic[TRequest, TResponse]):
     def __call__(
         self,
         request_iterator: typing.Iterator[TRequest],

diff --git a/typesafety/test_grpc.yml b/typesafety/test_grpc.yml
@@ -222,3 +222,29 @@
     server = typing.cast(aio.Server, None)
     enable_server_reflection(["foo"], server, None)
 
+- case: issue_41_stream_multi_callable_type_unary
+  main: |
+    import grpc
+    import typing
+    def unary_stream() -> grpc.UnaryStreamMultiCallable[int, int]:
+      return typing.cast(grpc.UnaryStreamMultiCallable[int, int], None)
+
+    stream = unary_stream()
+    stream.initial_metadata()
+    if stream.done() and stream.exception() is not None:
+      ex = stream.exception()
+      reveal_type(ex)  # N: Revealed type is "Union[builtins.Exception, None]"
+
+- case: issue_41_stream_multi_callable_type_stream
+  main: |
+    import grpc
+    import typing
+    def stream_stream() -> grpc.StreamStreamMultiCallable[int, int]:
+      return typing.cast(grpc.StreamStreamMultiCallable[int, int], None)
+
+    stream = stream_stream()
+    stream.initial_metadata()
+    if stream.done() and stream.exception() is not None:
+      ex = stream.exception()
+      reveal_type(ex)  # N: Revealed type is "Union[builtins.Exception, None]"
+