Parcourir la source

fix: stream_output argument in call_protobuf_handler

Denis Mazur il y a 4 ans
Parent
commit
56ade260dc
2 fichiers modifiés avec 11 ajouts et 3 suppressions
  1. 4 2
      hivemind/p2p/p2p_daemon.py
  2. 7 1
      hivemind/p2p/servicer.py

+ 4 - 2
hivemind/p2p/p2p_daemon.py

@@ -402,7 +402,7 @@ class P2P:
         input_protobuf_type: type,
         *,
         stream_input: bool = False,
-        stream_output:bool = False,
+        stream_output: bool = False,
     ) -> None:
         """
         :param stream_input: If True, assume ``handler`` to take ``TInputStream``
@@ -452,9 +452,11 @@ class P2P:
         name: str,
         input: Union[TInputProtobuf, TInputStream],
         output_protobuf_type: type,
+        *,
+        stream_output: bool = False,
     ) -> Awaitable[TOutputProtobuf]:
 
-        if not isinstance(input, AsyncIterableABC):
+        if not (isinstance(input, AsyncIterableABC) or stream_output):
             return await self._call_unary_protobuf_handler(peer_id, name, input, output_protobuf_type)
 
         responses = self._iterate_protobuf_stream_handler(peer_id, name, input, output_protobuf_type)

+ 7 - 1
hivemind/p2p/servicer.py

@@ -98,7 +98,13 @@ class ServicerBase:
                 self: StubBase, input: input_type, timeout: Optional[float] = None
             ) -> handler.response_type:
                 return await asyncio.wait_for(
-                    self._p2p.call_protobuf_handler(self._peer, handler.handle_name, input, handler.response_type),
+                    self._p2p.call_protobuf_handler(
+                        self._peer,
+                        handler.handle_name,
+                        input,
+                        handler.response_type,
+                        stream_output=handler.stream_output,
+                    ),
                     timeout=timeout,
                 )