From 2fdb8fc25a5de02c61f58f619858668a46a342e9 Mon Sep 17 00:00:00 2001
From: Tyler Goodlet <jgbt@protonmail.com>
Date: Mon, 16 Oct 2023 15:35:16 -0400
Subject: [PATCH] Factor non-yield stream msg processing into helper

Since both `MsgStream.receive()` and `.receive_nowait()` need the same
raising logic when a non-stream msg arrives (so that maybe an
appropriate IPC translated error can be raised) move the `KeyError`
handler code into a new `._streaming._raise_from_no_yield_msg()` func
and call it from both methods to make the error-interface-raising
symmetrical across both methods.
---
 tractor/_streaming.py | 112 +++++++++++++++++++++++++++---------------
 1 file changed, 72 insertions(+), 40 deletions(-)

diff --git a/tractor/_streaming.py b/tractor/_streaming.py
index 3045b83..e449fef 100644
--- a/tractor/_streaming.py
+++ b/tractor/_streaming.py
@@ -54,6 +54,60 @@ log = get_logger(__name__)
 #   messages? class ReceiveChannel(AsyncResource, Generic[ReceiveType]):
 # - use __slots__ on ``Context``?
 
+def _raise_from_no_yield_msg(
+    stream: MsgStream,
+    msg: dict,
+    src_err: KeyError,
+
+) -> bool:
+    '''
+    Raise an appopriate local error when a `MsgStream` msg arrives
+    which does not contain the expected (under normal operation)
+    `'yield'` field.
+
+    '''
+    # internal error should never get here
+    assert msg.get('cid'), ("Received internal error at portal?")
+
+    # TODO: handle 2 cases with 3.10+ match syntax
+    # - 'stop'
+    # - 'error'
+    # possibly just handle msg['stop'] here!
+
+    if stream._closed:
+        raise trio.ClosedResourceError('This stream was closed')
+
+    if msg.get('stop') or stream._eoc:
+        log.debug(f"{stream} was stopped at remote end")
+
+        # XXX: important to set so that a new ``.receive()``
+        # call (likely by another task using a broadcast receiver)
+        # doesn't accidentally pull the ``return`` message
+        # value out of the underlying feed mem chan!
+        stream._eoc = True
+
+        # # when the send is closed we assume the stream has
+        # # terminated and signal this local iterator to stop
+        # await stream.aclose()
+
+        # XXX: this causes ``ReceiveChannel.__anext__()`` to
+        # raise a ``StopAsyncIteration`` **and** in our catch
+        # block below it will trigger ``.aclose()``.
+        raise trio.EndOfChannel from src_err
+
+    # TODO: test that shows stream raising an expected error!!!
+    elif msg.get('error'):
+        # raise the error message
+        raise unpack_error(msg, stream._ctx.chan)
+
+    # always re-raise the source error if no translation error
+    # case is activated above.
+    raise src_err
+    # raise RuntimeError(
+    #     'Unknown non-yield stream msg?\n'
+    #     f'{msg}'
+    # )
+
 
 class MsgStream(trio.abc.Channel):
     '''
@@ -91,11 +145,20 @@ class MsgStream(trio.abc.Channel):
     # delegate directly to underlying mem channel
     def receive_nowait(self):
         msg = self._rx_chan.receive_nowait()
-        return msg['yield']
+        try:
+            return msg['yield']
+        except KeyError as kerr:
+            _raise_from_no_yield_msg(
+                stream=self,
+                msg=msg,
+                src_err=kerr,
+            )
 
     async def receive(self):
-        '''Async receive a single msg from the IPC transport, the next
-        in sequence for this stream.
+        '''
+        Receive a single msg from the IPC transport, the next in
+        sequence sent by the far end task (possibly in order as
+        determined by the underlying protocol).
 
         '''
         # see ``.aclose()`` for notes on the old behaviour prior to
@@ -110,43 +173,12 @@ class MsgStream(trio.abc.Channel):
             msg = await self._rx_chan.receive()
             return msg['yield']
 
-        except KeyError as err:
-            # internal error should never get here
-            assert msg.get('cid'), ("Received internal error at portal?")
-
-            # TODO: handle 2 cases with 3.10 match syntax
-            # - 'stop'
-            # - 'error'
-            # possibly just handle msg['stop'] here!
-
-            if self._closed:
-                raise trio.ClosedResourceError('This stream was closed')
-
-            if msg.get('stop') or self._eoc:
-                log.debug(f"{self} was stopped at remote end")
-
-                # XXX: important to set so that a new ``.receive()``
-                # call (likely by another task using a broadcast receiver)
-                # doesn't accidentally pull the ``return`` message
-                # value out of the underlying feed mem chan!
-                self._eoc = True
-
-                # # when the send is closed we assume the stream has
-                # # terminated and signal this local iterator to stop
-                # await self.aclose()
-
-                # XXX: this causes ``ReceiveChannel.__anext__()`` to
-                # raise a ``StopAsyncIteration`` **and** in our catch
-                # block below it will trigger ``.aclose()``.
-                raise trio.EndOfChannel from err
-
-            # TODO: test that shows stream raising an expected error!!!
-            elif msg.get('error'):
-                # raise the error message
-                raise unpack_error(msg, self._ctx.chan)
-
-            else:
-                raise
+        except KeyError as kerr:
+            _raise_from_no_yield_msg(
+                stream=self,
+                msg=msg,
+                src_err=kerr,
+            )
 
         except (
             trio.ClosedResourceError,  # by self._rx_chan