micropython/streampair: Package to create bi-directional linked stream objects.

micropython/streampair/manifest.py

@@ -0,0 +1,5 @@
+metadata(
+    description="Create a bi-directional linked pair of stream objects", version="0.0.1"
+)
+
+module("streampair.py")

micropython/streampair/streampair.py

@@ -0,0 +1,82 @@
+import io
+
+from micropython import RingIO, const
+
+try:
+    from typing import Union, Tuple
+except:
+    pass
+
+# From micropython/py/stream.h
+_MP_STREAM_ERROR = const(-1)
+_MP_STREAM_FLUSH = const(1)
+_MP_STREAM_SEEK = const(2)
+_MP_STREAM_POLL = const(3)
+_MP_STREAM_CLOSE = const(4)
+_MP_STREAM_POLL_RD = const(0x0001)
+
+
+def streampair(buffer_size: Union[int, Tuple[int, int]]=256):
+    """
+    Returns two bi-directional linked stream objects where writes to one can be read from the other and vice/versa.
+    This can be used somewhat similarly to a socket.socketpair in python, like a pipe
+    of data that can be used to connect stream consumers (eg. asyncio.StreamWriter, mock Uart)
+    """
+    try:
+        size_a, size_b = buffer_size
+    except TypeError:
+        size_a = size_b = buffer_size
+
+    a = RingIO(size_a)
+    b = RingIO(size_b)
+    return StreamPair(a, b), StreamPair(b, a)
+
+
+class StreamPair(io.IOBase):
+
+    def __init__(self, own: RingIO, other: RingIO):
+        self.own = own
+        self.other = other
+        super().__init__()
+
+    def read(self, nbytes=-1):
+        return self.own.read(nbytes)
+
+    def readline(self):
+        return self.own.readline()
+
+    def readinto(self, buf, limit=-1):
+        return self.own.readinto(buf, limit)
+
+    def write(self, data):
+        return self.other.write(data)
+
+    def seek(self, offset, whence):
+        return self.own.seek(offset, whence)
+
+    def flush(self):
+        self.own.flush()
+        self.other.flush()
+
+    def close(self):
+        self.own.close()
+        self.other.close()
+
+    def any(self):
+        return self.own.any()
+
+    def ioctl(self, op, arg):
+        if op == _MP_STREAM_POLL:
+            if self.any():
+                return _MP_STREAM_POLL_RD
+            return 0
+
+        elif op ==_MP_STREAM_FLUSH:
+            return self.flush()
+        elif op ==_MP_STREAM_SEEK:
+            return self.seek(arg[0], arg[1])
+        elif op ==_MP_STREAM_CLOSE:
+            return self.close()
+
+        else:
+            return _MP_STREAM_ERROR

micropython/streampair/test_streampair.py

@@ -0,0 +1,108 @@
+import asyncio
+import unittest
+import select
+from streampair import streampair
+
+def async_test(f):
+    """
+    Decorator to run an async test function
+    """
+    def wrapper(*args, **kwargs):
+        loop = asyncio.new_event_loop()
+        # loop.set_exception_handler(_exception_handler)
+        t = loop.create_task(f(*args, **kwargs))
+        loop.run_until_complete(t)
+
+    return wrapper
+
+
+class StreamPairTestCase(unittest.TestCase):
+    def test_streampair(self):
+        a, b = streampair()
+        assert a.write(b"foo") == 3
+        assert b.write(b"bar") == 3
+
+        assert (r := a.read()) == b"bar", r
+        assert (r := b.read()) == b"foo", r
+
+    @async_test
+    async def test_async_streampair(self):
+        a, b = streampair()
+        ar = asyncio.StreamReader(a)
+        bw = asyncio.StreamWriter(b)
+
+        br = asyncio.StreamReader(b)
+        aw = asyncio.StreamWriter(a)
+
+        aw.write(b"foo\n")
+        await aw.drain()
+        assert not a.any()
+        assert b.any()
+        assert (r := await br.readline()) == b"foo\n", r
+        assert not b.any()
+        assert not a.any()
+
+        bw.write(b"bar\n")
+        await bw.drain()
+        assert not b.any()
+        assert a.any()
+        assert (r := await ar.readline()) == b"bar\n", r
+        assert not b.any()
+        assert not a.any()
+
+    def test_select_poll_compatibility(self):
+        """Test that streampair works with select.poll()"""
+        a, b = streampair()
+
+        # Register stream with poll
+        poller = select.poll()
+        poller.register(a, select.POLLIN)
+
+        # No data available initially
+        events = poller.poll(0)
+        assert len(events) == 0, f"Expected no events, got {events}"
+
+        # Write data to b, should be readable from a
+        b.write(b"test data")
+
+        # Should now poll as readable
+        events = poller.poll(0)
+        assert len(events) == 1, f"Expected 1 event, got {events}"
+        assert events[0][0] == a, "Event should be for stream a"
+        assert events[0][1] & select.POLLIN, "Should be readable"
+
+        # Read the data
+        data = a.read()
+        assert data == b"test data", f"Expected b'test data', got {data}"
+
+        # Should no longer poll as readable
+        events = poller.poll(0)
+        assert len(events) == 0, f"Expected no events after read, got {events}"
+
+        poller.unregister(a)
+
+    @async_test
+    async def test_streamreader_direct_usage(self):
+        """Test that streampair can be used directly with asyncio.StreamReader"""
+        a, b = streampair()
+
+        # Create StreamReader directly on the streampair object
+        reader = asyncio.StreamReader(a)
+
+        # Write data in background task
+        async def write_delayed():
+            await asyncio.sleep_ms(10)
+            b.write(b"async test\n")
+
+        task = asyncio.create_task(write_delayed())
+
+        # Should be able to read via StreamReader
+        data = await asyncio.wait_for(reader.readline(), 1.0)
+        assert data == b"async test\n", f"Expected b'async test\\n', got {data}"
+
+        # Wait for background task to complete
+        await task
+
+
+if __name__ == "__main__":
+    unittest.main()