Set `Context._stream` in `Portal.open_stream_from()`..

runtime_to_msgspec
Tyler Goodlet 2024-04-24 12:43:08 -04:00
parent 18e97a8f9a
commit 6aa52417ef
1 changed files with 5 additions and 5 deletions

View File

@ -420,7 +420,6 @@ class Portal:
kwargs=kwargs, kwargs=kwargs,
portal=self, portal=self,
) )
ctx._portal = self
# ensure receive-only stream entrypoint # ensure receive-only stream entrypoint
assert ctx._remote_func_type == 'asyncgen' assert ctx._remote_func_type == 'asyncgen'
@ -430,9 +429,10 @@ class Portal:
async with MsgStream( async with MsgStream(
ctx=ctx, ctx=ctx,
rx_chan=ctx._rx_chan, rx_chan=ctx._rx_chan,
) as rchan: ) as stream:
self._streams.add(rchan) self._streams.add(stream)
yield rchan ctx._stream = stream
yield stream
finally: finally:
@ -454,7 +454,7 @@ class Portal:
# XXX: should this always be done? # XXX: should this always be done?
# await recv_chan.aclose() # await recv_chan.aclose()
self._streams.remove(rchan) self._streams.remove(stream)
# NOTE: impl is found in `._context`` mod to make # NOTE: impl is found in `._context`` mod to make
# reading/groking the details simpler code-org-wise. This # reading/groking the details simpler code-org-wise. This