moat-lib-stream 0.1.2__py3-none-any.whl

moat/lib/stream/__init__.py

@@ -0,0 +1,95 @@
+"""
+Stream infrastructure for handling data streams in a structured manner.
+"""
+
+from __future__ import annotations
+
+# Import base classes directly (not lazy)
+from .base import Base as Base
+from .base import BaseBlk as BaseBlk
+from .base import BaseBuf as BaseBuf
+from .base import BaseConn as BaseConn
+from .base import BaseMsg as BaseMsg
+from .base import StackedBlk as StackedBlk
+from .base import StackedBuf as StackedBuf
+from .base import StackedConn as StackedConn
+from .base import StackedMsg as StackedMsg
+
+from typing import TYPE_CHECKING as _TC
+
+if _TC:
+    from .anyio import AnyioBuf as AnyioBuf
+    from .anyio import BufAnyio as BufAnyio
+    from .anyio import FilenoBuf as FilenoBuf
+    from .anyio import ProcessBuf as ProcessBuf
+    from .anyio import ProcessDeadError as ProcessDeadError
+    from .anyio import RemoteBufAnyio as RemoteBufAnyio
+    from .anyio import SingleAnyioBuf as SingleAnyioBuf
+    from .cbor import CBORMsgBlk as CBORMsgBlk
+    from .cbor import CBORMsgBuf as CBORMsgBuf
+    from .log import LogBlk as LogBlk
+    from .log import LogBuf as LogBuf
+    from .log import LogMsg as LogMsg
+    from .reliable import EphemeralMsg as EphemeralMsg
+    from .reliable import ReliableMsg as ReliableMsg
+    from .tcp import TcpLink as TcpLink
+    from .terminal import FilenoTerm as FilenoTerm
+    from .terminal import TermBuf as TermBuf
+    from .unix import UnixLink as UnixLink
+
+
+# Lazy loading
+_imports = {
+    # Logging
+    "LogMsg": "log",
+    "LogBlk": "log",
+    "LogBuf": "log",
+    # CBOR
+    "CBORMsgBuf": "cbor",
+    "CBORMsgBlk": "cbor",
+    # AnyIO
+    "ProcessDeadError": "anyio",
+    "AnyioBuf": "anyio",
+    "FilenoBuf": "anyio",
+    "RemoteBufAnyio": "anyio",
+    "BufAnyio": "anyio",
+    "SingleAnyioBuf": "anyio",
+    "ProcessBuf": "anyio",
+    # Reliable messaging
+    "ReliableMsg": "reliable",
+    "EphemeralMsg": "reliable",
+    # Network connections
+    "TcpLink": "tcp",
+    "UnixLink": "unix",
+    # Terminal
+    "FilenoTerm": "terminal",
+    "TermBuf": "terminal",
+}
+
+
+def __getattr__(attr: str):
+    try:
+        mod = _imports[attr]
+    except KeyError:
+        raise AttributeError(attr) from None
+    value = getattr(__import__(mod, globals(), None, True, 1), attr)
+    globals()[attr] = value
+    return value
+
+
+__all__ = [
+    # Base classes (not lazy)
+    "Base",
+    "BaseBlk",
+    "BaseBuf",
+    "BaseConn",
+    "BaseMsg",
+    "StackedBlk",
+    "StackedBuf",
+    "StackedConn",
+    "StackedMsg",
+] + list(_imports.keys())
+
+
+def __dir__():
+    return __all__

moat/lib/stream/_cbor.py

@@ -0,0 +1,148 @@
+"""
+CBOR message encoding/decoding for stream layers.
+"""
+
+from __future__ import annotations
+
+from moat.lib.micro import AC_use, Lock, log
+from moat.lib.stream import BaseBuf, StackedMsg
+
+from ._console import _CReader
+
+from typing import TYPE_CHECKING  # isort:skip
+
+if TYPE_CHECKING:
+    from moat.lib.codec import Codec
+    from moat.util.liner import Liner
+
+    from collections.abc import Awaitable
+    from typing import Any
+
+
+class _CBORMsgBuf(StackedMsg):
+    """
+    structured messages > CBOR bytestream
+
+    Use this if your stream is reliable (TCP, USB, …) but doesn't support
+    message boundaries.
+
+    If @console is set and a prefix is used, sends data atomically.
+    Otherwise two separate write calls are used to save on message copying.
+
+    Config:
+        console (bool):
+            Flag how to handle non-framed data.
+            True: collect for crd/cwr, False: print incoming, None: ignore.
+        msg_prefix(int):
+            bytecode of prefix for messages (as opposed to console data)
+    """
+
+    cons = False
+    codec: Codec = None
+    liner: Liner = None
+
+    def __init__(self, stream: BaseBuf, cfg: dict):
+        #
+        # console: size of console buffer, 128 if True
+        # msg_prefix: int: code for start-of-packet
+        #
+        super().__init__(stream, cfg)
+        self.w_lock = Lock()
+
+        pref = cfg.get("msg_prefix")
+        if pref is not None:
+            pref = bytes((pref,))
+        self.pref = pref
+
+        cons = cfg.get("console", False)
+        if cons:
+            _CReader.__init__(self, cons)
+
+    async def setup(self):
+        await super().setup()
+        if self.cons is False:
+            from moat.util.liner import Liner  # noqa:PLC0415
+
+            self.liner = await AC_use(self, Liner())
+
+    async def cwr(self, buf):
+        if not self.cons:
+            return
+        return await self.s.wr(buf)
+
+    def crd(self, buf) -> Awaitable:
+        return _CReader.crd(self, buf)
+
+    async def send(self, msg: Any) -> None:
+        try:
+            msg = self.codec.encode(msg)
+        except Exception:
+            log("MSG:\n%r", msg)
+            raise
+        async with self.w_lock:
+            if self.pref is not None:
+                if True:  # self.cons:
+                    msg = self.pref + msg  # must be atomic
+                else:
+                    await self.s.wr(self.pref)
+            await self.s.wr(msg)
+
+    async def recv(self) -> Any:
+        """
+        Receive the next object.
+        """
+        # Pre+postcondition: the codec does not have an object in progress.
+
+        buf = bytearray(64)
+        if self.pref is None:
+            # easy case
+            while True:
+                try:
+                    r = next(self.codec)
+                except StopIteration:
+                    n = await self.s.rd(buf)
+                    self.codec.feed(memoryview(buf)[:n])
+                else:
+                    if self.cons and isinstance(r, int) and r >= 0:
+                        _CReader.cput(self, r)
+                    else:
+                        return r
+
+        while True:
+            b = bytearray(1)
+            # read until we get a prefix byte
+            if self.codec.unfeed(b) == 0:
+                n = await self.s.rd(buf)
+                self.codec.feed(memoryview(buf)[:n])
+            elif b == self.pref:
+                break
+            elif self.cons:
+                _CReader.cput(self, b[0])
+            elif self.liner is not None:
+                self.liner(b)
+
+        while True:
+            # read until we get an object
+            try:
+                return next(self.codec)
+            except StopIteration:
+                pass
+
+            n = await self.s.rd(buf)
+            self.codec.feed(memoryview(buf)[:n])
+
+
+class _CBORMsgBlk(StackedMsg):
+    """
+    structured messages > chunked bytestrings
+
+    Use this if the layer below supports byte boundaries
+    (one bytestring-ized message per call).
+    """
+
+    async def send(self, msg):
+        await self.s.snd(self.codec.encode(msg))
+
+    async def recv(self):
+        m = await self.s.rcv()
+        return self.codec.decode(m)

moat/lib/stream/_console.py

@@ -0,0 +1,62 @@
+"""
+Console data handling for stream layers.
+"""
+
+from __future__ import annotations
+
+from moat.lib.micro import Event
+
+
+class _CReader:
+    """
+    A mix-in that processes incoming console data.
+    """
+
+    def __init__(self, cons: bool | int):
+        if cons is True:
+            try:
+                import machine  # noqa: PLC0415,F401
+            except ImportError:
+                cons = 32768
+            else:
+                cons = 240
+        self.cevt = Event()
+        self.cpos = 0
+        self.cbuf = bytearray(cons)
+        self.cons = cons
+        self.intr = 0
+
+    async def crd(self, buf: bytearray):
+        """read waiting console data"""
+        if not self.cons:
+            raise EOFError
+        if not self.cpos:
+            await self.cevt.wait()
+            self.cevt = Event()
+        n = min(len(buf), self.cpos)
+        buf[:n] = self.cbuf[:n]
+        if n < self.cpos:
+            self.cbuf[: self.cpos - n] = self.cbuf[n : self.cpos]
+            self.cpos -= n
+        else:
+            self.cpos = 0
+        return n
+
+    def cput(self, b: int):
+        """store a byte in the console buffer"""
+        if self.cpos == len(self.cbuf):
+            if len(self.cbuf) > 100:
+                bfull = b"\n?BUF\n"
+                self.cbuf[0 : len(bfull)] = bfull
+                self.cpos = len(bfull)
+            else:
+                self.cpos = 0
+        if b != 3:
+            self.intr = 0
+        elif self.intr > 2:
+            raise KeyboardInterrupt
+        else:
+            self.intr += 1
+        self.cbuf[self.cpos] = b
+        self.cpos += 1
+        self.cevt.set()

moat/lib/stream/_serial.py

@@ -0,0 +1,72 @@
+"""
+SerialPacker protocol support for stream layers.
+"""
+
+from __future__ import annotations
+
+from moat.lib.micro import Lock
+from moat.lib.stream import BaseBuf, StackedBlk
+
+from ._console import _CReader
+
+
+class SerialPackerBlkBuf(StackedBlk):
+    """
+    chunked bytestrings > SerialPacker-ized stream
+
+    Use this (and a CBORHandler and a Reliable) if your AIO stream
+    is unreliable (TTL serial).
+    """
+
+    cons = False
+
+    def __init__(self, stream: BaseBuf, frame: dict, console: bool | int = False):
+        super().__init__(None)
+
+        from serialpacker import SerialPacker  # noqa: PLC0415
+
+        self.s = stream
+        self.p = SerialPacker(**frame)
+        self.buf = bytearray(16)
+        self.i = 0
+        self.n = 0
+        self.w_lock = Lock()
+        if console:
+            _CReader.__init__(self, console)
+
+    async def crd(self, buf):
+        if not self.cons:
+            raise EOFError
+        return await _CReader.crd(self, buf)
+
+    async def cwr(self, buf):
+        if not self.cons:
+            return
+        return await super().wr(buf)
+
+    async def recv(self):
+        while True:
+            while self.i < self.n:
+                msg = self.p.feed(self.buf[self.i])
+                self.i += 1
+                if isinstance(msg, int):
+                    if self.cons:
+                        _CReader.cput(self, msg)
+                elif msg is not None:
+                    return msg
+
+            n = await self.s.rd(self.buf)
+            if not n:
+                raise EOFError
+            self.i = 0
+            self.n = n
+
+    async def send(self, msg):
+        h, msg, t = self.p.frame(msg)
+        async with self.w_lock:
+            if not self.cons:
+                await self.s.wr(h)
+                await self.s.wr(msg)
+                await self.s.wr(t)
+            else:
+                await self.s.wr(h + msg + t)

moat/lib/stream/anyio.py

@@ -0,0 +1,285 @@
+"""
+AnyIO stream adaptors for MoaT (CPython-specific).
+"""
+
+from __future__ import annotations
+
+import anyio
+import fcntl
+import os
+import sys
+import termios
+from contextlib import asynccontextmanager
+
+from moat.util import CtxObj
+from moat.lib.micro import AC_use
+from moat.lib.stream import BaseBuf
+
+from typing import TYPE_CHECKING  # isort:skip
+
+if TYPE_CHECKING:
+    from moat.micro.cmd.tree.dir import SubDispatch
+
+
+class ProcessDeadError(RuntimeError):
+    """Process has died"""
+
+
+class AnyioBuf(BaseBuf):
+    """
+    Adapts an anyio stream to MoaT.
+    """
+
+    async def stream(self) -> anyio.abc.ByteStream:
+        """
+        Create the stream to use.
+
+        Use `AC_use` to arrange for closing it. This class will not do it
+        for you.
+        """
+        raise NotImplementedError(f"Override {self.__class__.__name__}.stream")
+
+    async def wr(self, buf) -> int:
+        "basic send"
+        try:
+            await self.s.send(buf)
+        except (anyio.EndOfStream, anyio.ClosedResourceError):
+            raise EOFError from None
+        else:
+            return len(buf)
+
+    async def rd(self, buf) -> int:
+        "basic receive-into"
+        try:
+            res = await self.s.receive(len(buf))
+        except (anyio.EndOfStream, anyio.ClosedResourceError):
+            raise EOFError from None
+        else:
+            buf[: len(res)] = res
+            return len(res)
+
+
+class FilenoBuf(BaseBuf):
+    """
+    Adapts a Unix file descriptor to MoaT.
+
+    """
+
+    rfd: int
+    wfd: int
+    rfl: int
+    wfl: int
+    term: bytes
+
+    def __init__(self, cfg, fd: int, wfd: int | None = None):
+        super().__init__(cfg)
+        self.rfd = fd
+        self.wfd = fd if wfd is None else wfd
+        self.term = None
+
+    async def stream(self) -> anyio.abc.ByteStream:
+        "Dummy here"
+        pass
+
+    async def setup(self):
+        "Change to nonblocking and raw"
+        self.rfl = fcntl.fcntl(self.rfd, fcntl.F_SETFL, 0)
+        if self.rfd != self.wfd:
+            self.wfl = fcntl.fcntl(self.wfd, fcntl.F_SETFL, 0)
+
+        fcntl.fcntl(self.rfd, fcntl.F_SETFL, self.rfl | os.O_NDELAY)
+        if self.rfd != self.wfd:
+            fcntl.fcntl(self.wfd, fcntl.F_SETFL, self.wfl | os.O_NDELAY)
+
+        try:
+            self.term = termios.tcgetattr(self.rfd)
+        except OSError:
+            pass
+        else:
+            new = self.term[:]
+            new[3] = new[3] & ~termios.ICANON & ~termios.ECHO & ~termios.ISIG
+            new[6][termios.VMIN] = 1
+            new[6][termios.VTIME] = 0
+            termios.tcsetattr(self.rfd, termios.TCSANOW, new)
+
+    async def teardown(self):
+        "Return to blocking and cooked"
+        fcntl.fcntl(self.rfd, fcntl.F_SETFL, self.rfl)
+        if self.rfd != self.wfd:
+            fcntl.fcntl(self.wfd, fcntl.F_SETFL, self.wfl)
+        if self.term is not None:
+            termios.tcsetattr(self.rfd, termios.TCSANOW, self.term)
+
+    async def wr(self, buf) -> int:
+        "basic send"
+        on = len(buf)
+        while n := len(buf):
+            await anyio.wait_writable(self.wfd)
+            nn = os.write(self.wfd, buf)
+            if nn <= 0:
+                raise OSError
+            if nn == n:
+                return on
+            buf = memoryview(buf)[nn:]
+
+    async def rd(self, buf) -> int:
+        "basic receive-into"
+        await anyio.wait_readable(self.rfd)
+        bf = os.read(self.rfd, len(buf))
+        buf[: len(bf)] = bf
+        return len(bf)
+
+
+class RemoteBufAnyio(anyio.abc.ByteStream):
+    """
+    Adapts a MoaT buf stream to a remote buffer read/write
+
+    TODO: use remote iteration for receiving
+    """
+
+    def __init__(self, disp: SubDispatch):
+        self.disp = disp
+
+    async def receive(self, max_bytes=256):
+        "forward to ``.rd``"
+        return await self.disp.rd(n=max_bytes)
+
+    async def send(self, buf):
+        "forward to ``.wr``"
+        await self.disp.wr(b=buf)
+
+    async def aclose(self):
+        "no-op"
+
+    async def send_eof(self):
+        "not implemented"
+        raise NotImplementedError("EOF")
+
+
+class BufAnyio(anyio.abc.ByteStream):
+    """
+    Adapts a MoaT Buf stream to an anyio bytestream.
+    """
+
+    par = None
+
+    def __init__(self, stream: BaseBuf):
+        self.stream = stream
+
+    async def __aenter__(self):
+        self.s = await self.stream.__aenter__()
+
+    async def __aexit__(self, *tb):
+        return await self.stream.__aexit__(*tb)
+
+    async def receive(self, max_bytes=256):
+        "forward to ``.rd``"
+        b = bytearray(max_bytes)
+        r = await self.s.rd(b)
+        if r == max_bytes:
+            return b
+        elif r <= max_bytes >> 2:
+            return bytes(b[:r])
+        else:
+            b = memoryview(b)
+            return b[:r]
+
+    async def send(self, buf):
+        "forward to ``.wr``"
+        await self.s.wr(buf)
+
+
+class SingleAnyioBuf(AnyioBuf):
+    """
+    Adapts an AnyIO stream to MoaT.
+
+    The stream is passed to the class constructor and can only be used
+    once.
+    """
+
+    def __init__(self, stream):
+        self._s = stream
+
+    async def stream(self):  # noqa:D102
+        return await AC_use(self, self._s)
+
+
+class ProcessBuf(CtxObj, AnyioBuf):
+    """
+    A stream that connects to an external process.
+
+    Config:
+    - exec: path to the executable
+    - argv: Arguments
+    - env: environment vars
+
+    You can set these as attributes on the object, statically or in your
+    subclass's `setup` method. Configuration can then override them.
+    """
+
+    proc: anyio.Process = None
+    exec: str | None = None
+    cwd: str | None = None
+    argv: list[str] | None = None
+    env: dict[str, str] | None = None
+
+    def __init__(self, cfg, executable: str | None = None, **kw):
+        super().__init__(cfg)
+        kw.setdefault("stderr", sys.stderr)
+        self.kw = kw
+        self.exec = executable
+
+    def open_args(self):
+        """Return keyword arguments for `anyio.open_process`.
+
+        Default is whatever has been passed to the ProcessBuf constructor.
+        """
+        # Ugh, anyio doesn't accept 'executable'
+        if self.exec is not None:
+            # self.kw["executable"] = self.exec
+            self.argv[0] = self.exec
+        elif "/" in (a0 := str(self.argv[0])):  # noqa:F841 # a0 unused
+            # self.kw["executable"] = a0
+            # self.argv[0] = a0.rsplit("/",1)[1]
+            pass
+        for k in ("cwd", "env"):
+            if (v := getattr(self, k)) is not None:
+                self.kw[k] = v
+
+        return self.kw
+
+    @asynccontextmanager
+    async def _ctx(self):
+        await self.setup()
+        proc = None
+        for k in ("exec", "argv", "env", "cwd"):
+            if (v := self.cfg.get(k, None)) is not None:
+                setattr(self, k, v)
+        if self.argv is None:
+            raise ValueError("Don't know what")
+
+        try:
+            async with await anyio.open_process(self.argv, **self.open_args()) as proc:
+                try:
+                    async with SingleAnyioBuf(
+                        anyio.streams.stapled.StapledByteStream(proc.stdin, proc.stdout),
+                    ) as s:
+                        yield s
+                    await proc.wait()
+                except BaseException:
+                    try:
+                        proc.kill()
+                        with anyio.CancelScope(shield=True):
+                            await proc.wait()
+                        raise
+                    finally:
+                        proc = None  # noqa:PLW2901
+        finally:
+            if proc is not None and proc.returncode != 0 and proc.returncode != -9:
+                raise ProcessDeadError(f"{self} died with {proc.returncode}")
+
+    async def setup(self):  # noqa:D102
+        pass
+
+    async def stream(self):  # noqa:D102
+        raise RuntimeError("should not be called")