forked from goodboy/tractor
1
0
Fork 0

Use `_raise_from_no_key_in_msg(allow_msgs)`

Instead of `allow_msg_keys` since we've fully flipped over to
struct-types for msgs in the runtime.

- drop the loop from `MsgStream.receive_nowait()` since
  `Yield/Return.pld` getting will handle both (instead of a loop of
  `dict`-key reads).
runtime_to_msgspec
Tyler Goodlet 2024-04-14 18:31:41 -04:00
parent f72b972348
commit 38a6483859
3 changed files with 19 additions and 33 deletions

View File

@ -2151,7 +2151,6 @@ async def open_context_from_portal(
src_err=src_error, src_err=src_error,
log=log, log=log,
expect_msg=Started, expect_msg=Started,
# expect_key='started',
) )
uid: tuple = portal.channel.uid uid: tuple = portal.channel.uid

View File

@ -938,7 +938,6 @@ def _raise_from_no_key_in_msg(
src_err: KeyError, src_err: KeyError,
log: StackLevelAdapter, # caller specific `log` obj log: StackLevelAdapter, # caller specific `log` obj
expect_key: str = 'yield',
expect_msg: str = Yield, expect_msg: str = Yield,
stream: MsgStream | None = None, stream: MsgStream | None = None,
@ -1053,7 +1052,7 @@ def _raise_from_no_key_in_msg(
# is activated above. # is activated above.
_type: str = 'Stream' if stream else 'Context' _type: str = 'Stream' if stream else 'Context'
raise MessagingError( raise MessagingError(
f"{_type} was expecting a '{expect_key.upper()}' message" f"{_type} was expecting a {expect_msg} message"
" BUT received a non-error msg:\n" " BUT received a non-error msg:\n"
f'{pformat(msg)}' f'{pformat(msg)}'
) from src_err ) from src_err

View File

@ -44,6 +44,7 @@ from .trionics import (
BroadcastReceiver, BroadcastReceiver,
) )
from tractor.msg import ( from tractor.msg import (
Return,
Stop, Stop,
Yield, Yield,
) )
@ -82,7 +83,7 @@ class MsgStream(trio.abc.Channel):
self, self,
ctx: Context, # typing: ignore # noqa ctx: Context, # typing: ignore # noqa
rx_chan: trio.MemoryReceiveChannel, rx_chan: trio.MemoryReceiveChannel,
_broadcaster: BroadcastReceiver | None = None, _broadcaster: BroadcastReceiver|None = None,
) -> None: ) -> None:
self._ctx = ctx self._ctx = ctx
@ -96,36 +97,26 @@ class MsgStream(trio.abc.Channel):
# delegate directly to underlying mem channel # delegate directly to underlying mem channel
def receive_nowait( def receive_nowait(
self, self,
allow_msg_keys: list[str] = ['yield'], allow_msgs: list[str] = Yield,
): ):
# msg: dict = self._rx_chan.receive_nowait()
msg: Yield|Stop = self._rx_chan.receive_nowait() msg: Yield|Stop = self._rx_chan.receive_nowait()
for ( # TODO: replace msg equiv of this or does the `.pld`
i, # interface read already satisfy it? I think so, yes?
key, try:
) in enumerate(allow_msg_keys): return msg.pld
try: except AttributeError as attrerr:
# return msg[key] _raise_from_no_key_in_msg(
return msg.pld ctx=self._ctx,
# except KeyError as kerr: msg=msg,
except AttributeError as attrerr: src_err=attrerr,
if i < (len(allow_msg_keys) - 1): log=log,
continue stream=self,
)
_raise_from_no_key_in_msg(
ctx=self._ctx,
msg=msg,
# src_err=kerr,
src_err=attrerr,
log=log,
expect_key=key,
stream=self,
)
async def receive( async def receive(
self, self,
hide_tb: bool = True, hide_tb: bool = False,
): ):
''' '''
Receive a single msg from the IPC transport, the next in Receive a single msg from the IPC transport, the next in
@ -157,10 +148,9 @@ class MsgStream(trio.abc.Channel):
try: try:
try: try:
msg: Yield = await self._rx_chan.receive() msg: Yield = await self._rx_chan.receive()
# return msg['yield']
return msg.pld return msg.pld
# except KeyError as kerr: # TODO: implement with match: instead?
except AttributeError as attrerr: except AttributeError as attrerr:
# src_err = kerr # src_err = kerr
src_err = attrerr src_err = attrerr
@ -170,10 +160,8 @@ class MsgStream(trio.abc.Channel):
_raise_from_no_key_in_msg( _raise_from_no_key_in_msg(
ctx=self._ctx, ctx=self._ctx,
msg=msg, msg=msg,
# src_err=kerr,
src_err=attrerr, src_err=attrerr,
log=log, log=log,
expect_key='yield',
stream=self, stream=self,
) )
@ -304,7 +292,7 @@ class MsgStream(trio.abc.Channel):
while not drained: while not drained:
try: try:
maybe_final_msg = self.receive_nowait( maybe_final_msg = self.receive_nowait(
allow_msg_keys=['yield', 'return'], allow_msgs=[Yield, Return],
) )
if maybe_final_msg: if maybe_final_msg:
log.debug( log.debug(