Denis Mazur 4 жил өмнө
parent
commit
71ca067fe7

+ 3 - 1
hivemind/p2p/p2p_daemon.py

@@ -402,13 +402,15 @@ class P2P:
         input_protobuf_type: type,
         *,
         stream_input: bool = False,
+        stream_output:bool = False,
     ) -> None:
         """
         :param stream_input: If True, assume ``handler`` to take ``TInputStream``
                              (not just ``TInputProtobuf``) as input.
+        :param stream_output: If True, assume ``handler`` to return ``TOutputStream``
         """
 
-        if not stream_input:
+        if not (stream_input or stream_output):
             await self._add_protobuf_unary_handler(name, handler, input_protobuf_type)
             return
 

+ 1 - 0
hivemind/p2p/servicer.py

@@ -113,6 +113,7 @@ class ServicerBase:
                 getattr(servicer, handler.method_name),
                 handler.request_type,
                 stream_input=handler.stream_input,
+                stream_output=handler.stream_output,
             )
 
     def get_stub(self, p2p: P2P, peer: PeerID) -> StubBase: