removed redundant context managers from websocket transport listeners

Author: gabi

rsocket/transports/aiohttp_websocket.py

@@ -52,11 +52,10 @@ async def connect(self):
         self._message_handler = asyncio.create_task(self.handle_incoming_ws_messages())
 
     async def handle_incoming_ws_messages(self):
-        with wrap_transport_exception():
-            async for msg in self._ws:
-                if msg.type == aiohttp.WSMsgType.BINARY:
-                    async for frame in self._frame_parser.receive_data(msg.data, 0):
-                        self._incoming_frame_queue.put_nowait(frame)
+        async for msg in self._ws:
+            if msg.type == aiohttp.WSMsgType.BINARY:
+                async for frame in self._frame_parser.receive_data(msg.data, 0):
+                    self._incoming_frame_queue.put_nowait(frame)
 
     async def send_frame(self, frame: Frame):
         with wrap_transport_exception():

rsocket/transports/aioquic_transport.py

@@ -83,15 +83,14 @@ async def send_frame(self, frame: Frame):
             await self._quic_protocol.query(frame)
 
     async def incoming_data_listener(self):
-        with wrap_transport_exception():
-            try:
-                while True:
-                    data = await self._incoming_bytes_queue.get()
-
-                    async for frame in self._frame_parser.receive_data(data, 0):
-                        self._incoming_frame_queue.put_nowait(frame)
-            except asyncio.CancelledError:
-                logger().debug('Asyncio task canceled: incoming_data_listener')
+        try:
+            while True:
+                data = await self._incoming_bytes_queue.get()
+
+                async for frame in self._frame_parser.receive_data(data, 0):
+                    self._incoming_frame_queue.put_nowait(frame)
+        except asyncio.CancelledError:
+            logger().debug('Asyncio task canceled: incoming_data_listener')
 
     async def close(self):
         self._listener.cancel()

rsocket/transports/quart_websocket.py

@@ -3,6 +3,7 @@
 from quart import websocket
 
 from rsocket.frame import Frame
+from rsocket.helpers import wrap_transport_exception
 from rsocket.logger import logger
 from rsocket.rsocket_server import RSocketServer
 from rsocket.transports.abstract_messaging import AbstractMessagingTransport
@@ -31,7 +32,8 @@ async def handle_incoming_ws_messages(self):
             logger().debug('Asyncio task canceled: quart_handle_incoming_ws_messages')
 
     async def send_frame(self, frame: Frame):
-        await websocket.send(frame.serialize())
+        with wrap_transport_exception():
+            await websocket.send(frame.serialize())
 
     async def close(self):
         pass