wish: improve web-stream Python connection object

wish/python/src/web_stream_ext.cc

@@ -385,11 +385,7 @@ NB_MODULE(web_stream_ext, m) {
           handler->SetOnMessage([&self](uint8_t opcode, const std::string& msg) {
             nb::gil_scoped_acquire acquire;
             try {
-              if (opcode == WEB_STREAM_OPCODE_BINARY || opcode == WEB_STREAM_OPCODE_METADATA) {
-                self.on_message_cb(opcode, nb::bytes(msg.data(), msg.size()));
-              } else {
-                self.on_message_cb(opcode, msg);
-              }
+              self.on_message_cb(opcode, nb::bytes(msg.data(), msg.size()));
             } catch (nb::python_error& e) {
               e.restore();
               PyErr_WriteUnraisable(self.on_message_cb.ptr());
@@ -501,11 +497,7 @@ NB_MODULE(web_stream_ext, m) {
           handler->SetOnMessage([&self](uint8_t opcode, const std::string& msg) {
             nb::gil_scoped_acquire acquire;
             try {
-              if (opcode == WEB_STREAM_OPCODE_BINARY || opcode == WEB_STREAM_OPCODE_METADATA) {
-                self.on_message_cb(opcode, nb::bytes(msg.data(), msg.size()));
-              } else {
-                self.on_message_cb(opcode, msg);
-              }
+              self.on_message_cb(opcode, nb::bytes(msg.data(), msg.size()));
             } catch (nb::python_error& e) {
               e.restore();
               PyErr_WriteUnraisable(self.on_message_cb.ptr());

wish/python/tests/test_e2e.py

@@ -129,6 +129,123 @@ async def run():
         asyncio.set_event_loop(loop)
         loop.run_until_complete(run())
 
+    def test_recv_decode_options(self):
+        async def run():
+            # Allow the server more time to start up to avoid Connection reset
+            await asyncio.sleep(1.0)
+
+            uri = f"webstream://127.0.0.1:{self.port}"
+
+            async with web_stream.connect(uri) as ws:
+                # 1. Test sending string and receiving with decode=False (returns bytes)
+                await ws.send("hello")
+                msg_bytes = await asyncio.wait_for(ws.recv(decode=False), timeout=2.0)
+                self.assertEqual(msg_bytes, b"hello")
+                self.assertIsInstance(msg_bytes, bytes)
+
+                # 2. Test sending string and receiving with decode=True (returns str)
+                await ws.send("world")
+                msg_str = await asyncio.wait_for(ws.recv(decode=True), timeout=2.0)
+                self.assertEqual(msg_str, "world")
+                self.assertIsInstance(msg_str, str)
+
+                # 3. Test sending bytes (which will be echoed as binary) and receiving with decode=True (returns str)
+                await ws.send(b"binary-text")
+                msg_str_from_bin = await asyncio.wait_for(ws.recv(decode=True), timeout=2.0)
+                self.assertEqual(msg_str_from_bin, "binary-text")
+                self.assertIsInstance(msg_str_from_bin, str)
+
+                # 4. Test sending bytes and receiving with default decode=None (returns bytes for binary)
+                await ws.send(b"binary-bytes")
+                msg_bytes_from_bin = await asyncio.wait_for(ws.recv(decode=None), timeout=2.0)
+                self.assertEqual(msg_bytes_from_bin, b"binary-bytes")
+                self.assertIsInstance(msg_bytes_from_bin, bytes)
+
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+        loop.run_until_complete(run())
+
+    def test_recv_metadata(self):
+        async def run():
+            # Allow the server more time to start up to avoid Connection reset
+            await asyncio.sleep(1.0)
+
+            uri = f"webstream://127.0.0.1:{self.port}"
+
+            async with web_stream.connect(uri) as ws:
+                # We can manually put a metadata message into ws._recv_queue
+                await ws._recv_queue.put((web_stream.WEB_STREAM_OPCODE_METADATA, b"meta-data-val"))
+
+                # Default decode=None: should return bytes
+                msg = await ws.recv(decode=None)
+                self.assertEqual(msg, b"meta-data-val")
+                self.assertIsInstance(msg, bytes)
+
+                # Put another one and test decode=True: should return str
+                await ws._recv_queue.put((web_stream.WEB_STREAM_OPCODE_METADATA, b"meta-data-val2"))
+                msg_str = await ws.recv(decode=True)
+                self.assertEqual(msg_str, "meta-data-val2")
+                self.assertIsInstance(msg_str, str)
+
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+        loop.run_until_complete(run())
+
+    def test_iterator_break(self):
+        async def run():
+            # Allow the server more time to start up to avoid Connection reset
+            await asyncio.sleep(1.0)
+
+            uri = f"webstream://127.0.0.1:{self.port}"
+
+            async with web_stream.connect(uri) as ws:
+                await ws.send("hello")
+                await ws.send("world")
+
+                received = []
+                async for msg in ws:
+                    received.append(msg)
+                    if len(received) == 2:
+                        break
+
+                self.assertEqual(received, ["hello", "world"])
+
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+        try:
+            loop.run_until_complete(run())
+        finally:
+            loop.run_until_complete(loop.shutdown_asyncgens())
+            loop.close()
+
+    def test_iterator_close(self):
+        async def run():
+            # Allow the server more time to start up to avoid Connection reset
+            await asyncio.sleep(1.0)
+
+            uri = f"webstream://127.0.0.1:{self.port}"
+
+            async with web_stream.connect(uri) as ws:
+                async def close_later():
+                    await asyncio.sleep(0.5)
+                    await ws.close()
+
+                asyncio.create_task(close_later())
+
+                messages = []
+                async for msg in ws:
+                    messages.append(msg)
+
+                self.assertEqual(messages, [])
+
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+        try:
+            loop.run_until_complete(run())
+        finally:
+            loop.run_until_complete(loop.shutdown_asyncgens())
+            loop.close()
+
     def test_server_disconnect_detect(self):
         port = get_free_port()
         cmd = [

wish/python/web_stream/__init__.py

@@ -1,3 +1,15 @@
-from .client import connect, WebStreamConnection
+from .client import (
+    connect,
+    WebStreamConnection,
+    WEB_STREAM_OPCODE_TEXT,
+    WEB_STREAM_OPCODE_BINARY,
+    WEB_STREAM_OPCODE_METADATA,
+)
 
-__all__ = ["connect", "WebStreamConnection"]
+__all__ = [
+    "connect",
+    "WebStreamConnection",
+    "WEB_STREAM_OPCODE_TEXT",
+    "WEB_STREAM_OPCODE_BINARY",
+    "WEB_STREAM_OPCODE_METADATA",
+]

wish/python/web_stream/client.py

@@ -4,6 +4,11 @@
 
 from . import web_stream_ext
 
+# WebStream Opcodes
+WEB_STREAM_OPCODE_TEXT = 1
+WEB_STREAM_OPCODE_BINARY = 2
+WEB_STREAM_OPCODE_METADATA = 3
+
 class WebStreamConnection:
     def __init__(self, host, port, tls, ca_file="", cert_file="", key_file=""):
         self._host = host
@@ -74,7 +79,7 @@ async def send(self, data):
         else:
             self._handler.send_text(str(data))
 
-    async def recv(self):
+    async def recv(self, decode=None):
         """Receives a message from the WebStream connection."""
         if not self._client:
             raise RuntimeError("Connection is closed")
@@ -85,14 +90,32 @@ async def recv(self):
             raise res
 
         opcode, msg = res
-        # You can process opcode here if you want to distinguish text/binary
-        # We'll just return the message
-        # In actual implementation: 1=Text, 2=Binary
-        if opcode == 2:
-            return msg.encode("utf-8") if isinstance(msg, str) else msg
+
+        # Determine if we should return str or bytes
+        # Default behavior:
+        # - WEB_STREAM_OPCODE_TEXT -> str
+        # - WEB_STREAM_OPCODE_BINARY -> bytes
+        # - WEB_STREAM_OPCODE_METADATA -> bytes
+        # Note: UTF-8 decoding for metadata is undefined in the specification.
+        # We treat it similarly to binary messages for now (i.e. no UTF-8 decoding by default).
+        should_decode = decode
+        if should_decode is None:
+            should_decode = (opcode == WEB_STREAM_OPCODE_TEXT)
+
+        if should_decode:
+            return msg.decode("utf-8")
         else:
             return msg
 
+    async def __aiter__(self):
+        """Iterate on incoming messages."""
+        try:
+            while True:
+                yield await self.recv()
+        except (ConnectionAbortedError, RuntimeError):
+            pass
+
+
 class _ConnectContextManager:
     def __init__(self, uri, ca_file="", cert_file="", key_file=""):
         parsed = urlparse(uri)