Skip to content

Commit 80d2075

Browse files
committed
Add another interface for ClientStubStreamStream
1 parent ede443e commit 80d2075

File tree

1 file changed

+14
-4
lines changed

1 file changed

+14
-4
lines changed

src/purerpc/wrappers.py

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -88,8 +88,18 @@ async def __call__(self, message_aiter):
8888

8989

9090
class ClientStubStreamStream(ClientStub):
91-
async def __call__(self, message_aiter=None):
91+
async def call_aiter(self, message_aiter):
9292
stream = await self._stream_fn()
93-
await curio.spawn(send_multiple_messages_finalize, stream, message_aiter, daemon=True)
94-
async for message in stream_to_async_iterator(stream):
95-
yield message
93+
if message_aiter is not None:
94+
await curio.spawn(send_multiple_messages_finalize, stream, message_aiter, daemon=True)
95+
async for message in stream_to_async_iterator(stream):
96+
yield message
97+
98+
async def call_stream(self):
99+
return await self._stream_fn()
100+
101+
def __call__(self, message_aiter=None):
102+
if message_aiter is None:
103+
return self.call_stream()
104+
else:
105+
return self.call_aiter(message_aiter)

0 commit comments

Comments
 (0)