socketwrapper 0.0.1__py3-none-any.whl

socketwrapper/__init__.py

@@ -0,0 +1,476 @@
+"""socketwrapper - Message socket wrappers."""
+
+import asyncio
+import collections.abc
+import os
+import selectors
+import socket
+import stat
+import sys
+import typing
+
+
+if sys.version_info < (3, 13):
+    import typing_extensions
+else:
+    typing_extensions = typing
+
+
+from . import _utils, framing, protocols
+
+
+__all__ = (
+    # type aliases and protocols
+    'MessageFraming',
+    'RecvSize',
+    'SendPayload',
+    'SocketLike',
+    'SocketOrDescriptor',
+
+    # wrappers
+    'SocketWrapper',
+    'SocketReader',
+    'SocketWriter',
+    'SocketDuplex',
+
+    'MessageWrapper',
+    'MessageReader',
+    'MessageWriter',
+    'MessageDuplex',
+
+    # utils
+    'pipe',
+    'socketpair',
+    )
+
+
+DEFAULT_SOCKETPAIR_FAMILY = getattr(socket, 'AF_UNIX', socket.AF_INET)
+DEFAULT_FRAMING = framing.VarIntBytes
+
+
+SocketLike = protocols.SocketLike
+MessageFraming = protocols.MessageFraming
+SocketOrDescriptor = SocketLike | protocols.HasFileno | int
+RecvSize = _utils.RecvSize
+SendPayload = _utils.SendPayload
+
+
+class DescriptorSocket:
+    """Generic socket-like accepting any file descriptor (usualy pipes) as input."""
+
+    def __init__(self, fd: int) -> None:
+        """Initialize."""
+        self._fd = fd
+
+    def fileno(self) -> int:
+        """Get underlying file descriptor."""
+        return self._fd
+
+    def send(self, data: collections.abc.Buffer) -> int:
+        """Write data to descriptor."""
+        return os.write(self._fd, data)
+
+    def recv(self, bufsize: int) -> bytes:
+        """Read data from descriptor."""
+        return os.read(self._fd, bufsize)
+
+    def settimeout(self, timeout: float | None) -> None:
+        """Set blocking or non-blocking based on timeout being 0 or None."""
+        if timeout:
+            raise ValueError(timeout)
+
+        os.set_blocking(self._fd, timeout is None)
+
+    def close(self) -> None:
+        """Close file descriptor."""
+        os.close(self._fd)
+
+
+class SocketWrapper(_utils.ClosingContext):
+    """Base class for socket wrappers supporting asynchronous operations."""
+
+    _sock: SocketLike
+    _duplex: bool = False
+    _wselector: selectors.BaseSelector | None = None
+    _rselector: selectors.BaseSelector | None = None
+
+    def __init__(self, sock: SocketOrDescriptor) -> None:
+        """Initialize for socket."""
+        self._sock = _socketlike(sock)
+
+    @property
+    def inheritable(self) -> bool:
+        """Get whether or not this socket is inheritable by subprocesses."""
+        return os.get_inheritable(self._sock.fileno())
+
+    @inheritable.setter
+    def inheritable(self, inheritable: bool) -> None:
+        """Set whether or not this socket is inheritable by subprocesses."""
+        os.set_inheritable(self._sock.fileno(), inheritable)
+
+    def fileno(self) -> int:
+        """Get underlying socket file descriptor."""
+        return self._sock.fileno()
+
+    def close(self) -> None:
+        """Close underlying socket."""
+        for obj in (self._rselector, self._wselector, self._sock):
+            if obj:
+                obj.close()
+
+    def _selector(self, *, write: bool = False) -> collections.abc.Callable[[float | None], list[tuple]]:
+        """Get selector for socket."""
+        sel = self._wselector if write else self._rselector
+        if not sel:
+            sel = selectors.DefaultSelector()
+            if write:
+                sel.register(self._sock.fileno(), selectors.EVENT_WRITE)
+                self._wselector = sel
+
+            else:
+                sel.register(self._sock.fileno(), selectors.EVENT_READ)
+                self._rselector = sel
+
+        return sel.select
+
+    def _consume(self, processor: collections.abc.Iterable, deadline: float | None, *, write: bool = False) -> None:
+        """Configure socket and consume processor."""
+        timeout = _utils.timeout_checker(deadline, 'timeout during transmission') if deadline else None
+        wait = self._sock.settimeout if timeout else None
+        if self._duplex or _utils.capture(ValueError, self._sock.settimeout, timeout() if timeout else None):
+            self._sock.settimeout(.0)
+            wait = self._selector(write=write)
+
+        if wait and timeout:
+            for _ in processor:
+                wait(timeout())
+
+        elif wait:
+            for _ in processor:
+                wait(None)
+
+        else:
+            for _ in processor:
+                pass
+
+    async def _consume_async(self, processor: collections.abc.Iterable, *, write: bool = False) -> None:
+        """Configure socket as non-blocking and consume processor."""
+
+        def notify() -> None:
+            """Pass buffer to target and update future (and clean up) on success or error."""
+            try:
+                if next(iterator, True):
+                    future.set_result(None)
+                    remove_io(self)
+
+            except asyncio.InvalidStateError:
+                remove_io(self)
+
+            except Exception as e:
+                future.set_exception(e)
+                remove_io(self)
+
+        loop = asyncio.get_event_loop()
+        future = loop.create_future()
+        iterator = iter(processor)
+        add_io, remove_io = (
+            (loop.add_writer, loop.remove_writer) if write else
+            (loop.add_reader, loop.remove_reader)
+            )
+
+        self._sock.settimeout(.0)
+        add_io(self, notify)
+        try:
+            await future
+
+        finally:
+            remove_io(self)
+
+
+class SocketReader(SocketWrapper):
+    """Readable socket wrapper."""
+
+    _rlock: _utils.CrossLock
+
+    def __init__(self, sock: SocketOrDescriptor) -> None:
+        """Initialize for socket."""
+        super().__init__(sock)
+        self._rlock = _utils.CrossLock()
+
+    def _recv(self, bufsize: RecvSize, timeout: float | None = None) -> bytearray:
+        """Receive data from socket as bytearray."""
+        with _utils.lock_timeout(self._rlock, timeout=timeout) as deadline:
+            buffer = bytearray()
+            reader = _utils.reader(self._sock, buffer, bufsize)
+            self._consume(reader, deadline)
+            return buffer
+
+    async def _recv_async(self, bufsize: RecvSize) -> bytearray:
+        """Receive data from socket (async) as bytearray."""
+        async with self._rlock:
+            buffer = bytearray()
+            reader = _utils.reader(self._sock, buffer, bufsize)
+            await self._consume_async(reader)
+            return buffer
+
+    def recv(self, bufsize: int, timeout: float | None = None) -> bytes:
+        """Receive data from socket as bytes."""
+        return bytes(self._recv(bufsize, timeout))
+
+    async def recv_async(self, bufsize: int) -> bytes:
+        """Receive data from socket (async) as bytes."""
+        return bytes(await self._recv_async(bufsize))
+
+
+class SocketWriter(SocketWrapper):
+    """Writable socket wrapper."""
+
+    _wlock: _utils.CrossLock
+
+    def __init__(self, sock: SocketOrDescriptor) -> None:
+        """Initialize for socket."""
+        super().__init__(sock)
+        self._wlock = _utils.CrossLock()
+
+    def send(self, data: SendPayload, timeout: float | None = None) -> None:
+        """Send data to socket."""
+        with _utils.lock_timeout(self._wlock, timeout=timeout) as deadline:
+            writer = _utils.writer(self._sock, data)
+            self._consume(writer, deadline, write=True)
+
+    async def send_async(self, data: SendPayload) -> None:
+        """Send data to socket (async)."""
+        async with self._wlock:
+            writer = _utils.writer(self._sock, data)
+            await self._consume_async(writer, write=True)
+
+
+class SocketDuplex(SocketWriter, SocketReader):
+    """Duplex (both readable and writable) socket wrapper."""
+
+    _duplex: bool = True
+
+
+# TODO(py313+): use new syntax
+R = typing.TypeVar('R')
+W = typing.TypeVar('W')
+KS = typing_extensions.TypeVar('KS', bound=SocketWrapper, default=SocketWrapper)
+KSR = typing_extensions.TypeVar('KSR', bound=SocketReader, default=SocketReader)
+KSW = typing_extensions.TypeVar('KSW', bound=SocketWriter, default=SocketWriter)
+
+
+class MessageWrapper(_utils.ClosingContext, typing_extensions.Generic[R, W, KS]):
+    """Base class for message socket wrappers supporting asynchronous operations."""
+
+    _wrapper: typing.ClassVar[type[SocketWrapper]] = SocketWrapper
+    _raw: KS
+    framing: MessageFraming[R, W]
+
+    def __init__(self, sock: KS | SocketOrDescriptor, framing: MessageFraming[R, W] = DEFAULT_FRAMING) -> None:
+        """Initialize for socket."""
+        self._raw = sock if isinstance(sock, SocketWrapper) else self._wrapper(sock)
+        self.framing = framing
+
+    @property
+    def inheritable(self) -> bool:
+        """Get whether or not this socket is inheritable by subprocesses."""
+        return self._raw.inheritable
+
+    @inheritable.setter
+    def inheritable(self, inheritable: bool) -> None:
+        """Set whether or not this socket is inheritable by subprocesses."""
+        self._raw.inheritable = inheritable
+
+    def fileno(self) -> int:
+        """Get underlying socket file descriptor."""
+        return self._raw.fileno()
+
+    def close(self) -> None:
+        """Close underlying socket."""
+        self._raw.close()
+
+
+class MessageReader(MessageWrapper[R, typing.Any, KSR], typing_extensions.Generic[R, KSR]):
+    """Readable message socket wrapper."""
+
+    _wrapper: typing.ClassVar = SocketReader
+
+    def recv(self, timeout: float | None = None) -> R:
+        """Receive data from socket."""
+        return self.framing.loads(self._raw._recv(self.framing.frames, timeout=timeout))
+
+    async def recv_async(self) -> R:
+        """Receive data from socket (async)."""
+        return self.framing.loads(await self._raw._recv_async(self.framing.frames))
+
+
+class MessageWriter(MessageWrapper[typing.Any, W, KSW], typing_extensions.Generic[W, KSW]):
+    """Writable message socket wrapper."""
+
+    _wrapper: typing.ClassVar = SocketWriter
+
+    def send(self, data: W, timeout: float | None = None) -> None:
+        """Send data to socket."""
+        self._raw.send(self.framing.dumps(data), timeout=timeout)
+
+    async def send_async(self, data: W) -> None:
+        """Send data to socket (async)."""
+        await self._raw.send_async(self.framing.dumps(data))
+
+
+class MessageDuplex[R, W](MessageWriter[W, SocketDuplex], MessageReader[R, SocketDuplex]):
+    """Duplex (both readable and writable) message socket wrapper."""
+
+    _wrapper: typing.ClassVar = SocketDuplex
+
+
+# TODO: move to namespace
+_WrapperPair = _utils.ContextPair[SocketReader, SocketWriter]
+_DuplexWrapperPair = _utils.ContextPair[SocketDuplex, SocketDuplex]
+_DefaultMessagePair = _utils.ContextPair[MessageReader[bytes], MessageWriter[protocols.SizedBuffer]]
+_DuplexDefaultMessagePair = _utils.ContextPair[
+    MessageDuplex[bytes, protocols.SizedBuffer],
+    MessageDuplex[bytes, protocols.SizedBuffer],
+    ]
+
+
+@typing.overload
+def socketpair(
+    family: int = ...,
+    type: typing.Literal[socket.SocketKind.SOCK_STREAM] = socket.SOCK_STREAM,
+    proto: int = ...,
+    framing: typing.Literal[False] = False,
+    ) -> _DuplexWrapperPair: ...
+
+
+@typing.overload
+def socketpair(
+    family: int = ...,
+    type: socket.SocketKind | int = socket.SOCK_STREAM,
+    proto: int = ...,
+    framing: typing.Literal[False] = False,
+    ) -> _WrapperPair: ...
+
+
+@typing.overload
+def socketpair(
+    family: int = ...,
+    type: typing.Literal[socket.SocketKind.SOCK_STREAM] = socket.SOCK_STREAM,
+    proto: int = ...,
+    framing: typing.Literal[True] = ...,
+    ) -> _DefaultMessagePair: ...
+
+
+@typing.overload
+def socketpair(
+    family: int = ...,
+    type: socket.SocketKind | int = ...,
+    proto: int = ...,
+    framing: typing.Literal[True] = ...,
+    ) -> _DuplexDefaultMessagePair: ...
+
+
+@typing.overload
+def socketpair[R, W](
+    family: int = ...,
+    type: typing.Literal[socket.SocketKind.SOCK_STREAM] = socket.SOCK_STREAM,
+    proto: int = ...,
+    framing: MessageFraming[R, W] = ...,
+    ) -> _utils.ContextPair[MessageDuplex[R, W], MessageDuplex[R, W]]: ...
+
+
+@typing.overload
+def socketpair[R, W](
+    family: int = ...,
+    type: socket.SocketKind | int = ...,
+    proto: int = ...,
+    framing: MessageFraming[R, W] = ...,
+    ) -> _utils.ContextPair[MessageReader[R], MessageWriter[W]]: ...
+
+
+def socketpair(
+        family=DEFAULT_SOCKETPAIR_FAMILY,
+        type=socket.SOCK_STREAM,  # noqa: A002
+        proto=0,
+        framing=False,
+        ):
+    """Initialize connected socket pair as socket wrappers, or message wrappers if framing is enabled."""
+    sock_a, sock_b = socket.socketpair(family, type, proto)
+    return (
+        _utils.ContextPair(_duplex(sock_a, framing), _duplex(sock_b, framing)) if type == socket.SOCK_STREAM else
+        _utils.ContextPair(_reader(sock_a, framing), _writer(sock_b, framing))
+        )
+
+
+@typing.overload
+def pipe(framing: typing.Literal[False] = False) -> _WrapperPair: ...
+
+
+@typing.overload
+def pipe(framing: typing.Literal[True]) -> _DefaultMessagePair: ...
+
+
+@typing.overload
+def pipe[R, W](framing: MessageFraming[R, W]) -> _utils.ContextPair[MessageReader[R], MessageWriter[W]]: ...
+
+
+def pipe(framing=False):
+    """Initialize pipe (readable, writable) as socket wrapper pair, or message wrappers pair if framing is enabled."""
+    read_fd, write_fd = os.pipe()
+    return _utils.ContextPair(_reader(read_fd, framing), _writer(write_fd, framing))
+
+
+def _reader[R, W](sock: SocketOrDescriptor, framing: MessageFraming[R, W] | bool) -> SocketReader | MessageReader[R]:
+    """Initialize bytes or message reader wrapper based on framing."""
+    return MessageReader(sock, fr) if (fr := _framing(framing)) else SocketReader(sock)
+
+
+def _writer[R, W](sock: SocketOrDescriptor, framing: MessageFraming[R, W] | bool) -> SocketWriter | MessageWriter[W]:
+    """Initialize bytes or message writer wrapper based on framing."""
+    return MessageWriter(sock, fr) if (fr := _framing(framing)) else SocketWriter(sock)
+
+
+def _duplex[R, W](sock: SocketOrDescriptor, framing: MessageFraming[R, W] | bool) -> SocketDuplex | MessageDuplex[R, W]:
+    """Initialize duplex bytes or message wrapper based on framing."""
+    return MessageDuplex(sock, fr) if (fr := _framing(framing)) else SocketDuplex(sock)
+
+
+# TODO(py313+): use new syntax
+DR = typing_extensions.TypeVar('DR', default=bytes)
+DW = typing_extensions.TypeVar('DW', default=protocols.SizedBuffer)
+
+
+def _framing(framing: MessageFraming[DR, DW] | bool) -> MessageFraming[DR, DW] | None:
+    """Pick framing (or default framing) based on framing parameter."""
+    return DEFAULT_FRAMING if framing is True else framing if framing else None
+
+
+def _socketlike(sock: SocketOrDescriptor) -> SocketLike:
+    """Use or initialize socketlike, reporting if it's native or not."""
+    if isinstance(sock, SocketLike):
+        return sock  # socket object, use straight
+
+    fd, dup = (sock, False) if isinstance(sock, int) else (sock.fileno(), True)
+
+    try:
+        maybe_socket = os.name == 'nt' or stat.S_ISSOCK(os.fstat(fd).st_mode)
+
+    except Exception:
+        maybe_socket = False
+
+    if maybe_socket:
+        sock = socket.socket(fileno=fd)
+        try:
+            sock.getsockopt(socket.SOL_SOCKET, socket.SO_TYPE)
+            if dup:  # obj backed by socket, duplicate to not interfere with original ref lifetime
+                sock2 = sock.dup()
+                sock.detach()
+                return sock2
+
+        except OSError:
+            sock.detach()
+
+        else:
+            return sock  # socket descriptor, use directly
+
+    return DescriptorSocket(fd)  # assume pipe-like

socketwrapper/_utils.py

@@ -0,0 +1,283 @@
+"""Utility module."""
+
+import abc
+import asyncio
+import collections.abc
+import contextlib
+import socket
+import threading
+import time
+import types
+import typing
+import weakref
+
+from . import protocols
+
+
+FrameFunction = collections.abc.Callable[[bytearray], collections.abc.Iterable[int]]
+RecvSize = FrameFunction | int
+SendPayload = collections.abc.Iterable[protocols.SizedBuffer] | protocols.SizedBuffer
+
+
+class ClosingContext(abc.ABC):
+
+    @abc.abstractmethod
+    def close(self) -> None:
+        """Perform close operations (abstract)."""
+
+    def __enter__(self) -> typing.Self:
+        """Get self class as context manager."""
+        return self
+
+    def __exit__(
+            self,
+            exc_type: type | None,
+            exc_value: BaseException | None,
+            exc_traceback: types.TracebackType | None,
+            ) -> None:
+        """Leave context manager closing the socket."""
+        self.close()
+
+
+class ContextPair[A: contextlib.AbstractContextManager, B: contextlib.AbstractContextManager](tuple[A, B]):
+    """Tuple supporting context manager protocol."""
+
+    def __new__(cls, a: A, b: B) -> typing.Self:
+        """Initialize with variadic arguments as items (instead of iterable)."""
+        return super().__new__(cls, (a, b))
+
+    def __enter__(self) -> typing.Self:
+        """Enter context on all items and return tuple itself."""
+        errors = []
+        for item in self:
+            try:
+                item.__enter__()
+
+            except Exception as err:
+                errors.append(err)
+
+        if errors:
+            message = 'Context manager exceptions'
+            raise ExceptionGroup(message, errors)
+
+        return self
+
+    def __exit__(
+            self,
+            exc_type: type | None,
+            exc_value: BaseException | None,
+            exc_traceback: types.TracebackType | None,
+            ) -> bool | None:
+        """Leave context on all items."""
+        booleans = 0
+        truthies = 0
+        errors = []
+        for item in self:
+            try:
+                res = item.__exit__(exc_type, exc_value, exc_traceback)
+                booleans += res is not None
+                truthies += bool(res)
+
+            except Exception as err:
+                errors.append(err)
+
+        if errors:
+            message = 'Context manager exceptions'
+            raise ExceptionGroup(message, errors)
+
+        return True if truthies else False if booleans else None
+
+
+class CrossLock:
+    """A higher level lock for both sync and async."""
+
+    __slots__: typing.ClassVar = '_aiolocks', '_aiowaits', '_busy', '_lock'
+
+    _aiolocks: weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, asyncio.Lock]
+    _aiowaits: weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, asyncio.Future]
+    _busy: threading.Lock
+    _lock: threading.Lock
+
+    def __init__(self) -> None:
+        """Initialize lock state."""
+        self._lock = threading.Lock()
+        self._busy = threading.Lock()
+        self._aiolocks = weakref.WeakKeyDictionary()
+        self._aiowaits = weakref.WeakKeyDictionary()
+
+    @property
+    def locked(self) -> bool:
+        return self._lock.locked()
+
+    def acquire(self, *, timeout: float | None = None) -> bool:
+        return (
+            self._lock.acquire() if timeout in {None, -1} else
+            self._lock.acquire(timeout=timeout) if timeout else
+            self._lock.acquire(blocking=False)
+            )
+
+    def release(self) -> None:
+        with self._busy:
+            aiowaits = dict(self._aiowaits)
+            self._aiowaits.clear()
+            self._lock.release()
+
+        try:
+            if future := aiowaits.pop(asyncio.get_running_loop(), None):
+                future.set_result(None)
+
+        except RuntimeError:
+            pass
+
+        for loop, future in aiowaits.items():
+            loop.call_soon_threadsafe(future.set_result, None)
+
+    async def acquire_async(self) -> typing.Literal[True]:
+        loop = asyncio.get_running_loop()
+        aiolock = self._aiolocks.get(loop) or self._aiolocks.setdefault(loop, asyncio.Lock())
+
+        async with aiolock:  # per-loop lock mitigating thundering herd
+            while True:
+                with self._busy:
+                    if self._lock.acquire(blocking=False):
+                        return True
+
+                    future = self._aiowaits.get(loop) or self._aiowaits.setdefault(loop, loop.create_future())
+
+                await asyncio.shield(future)  # wait for release notification
+
+    def __await__(self) -> collections.abc.Generator[None, None, None]:
+        yield from self.acquire_async().__await__()
+
+    def __enter__(self) -> typing.Self:
+        self.acquire()
+        return self
+
+    def __exit__(
+            self,
+            exc_type: type | None,
+            exc_value: BaseException | None,
+            exc_traceback: types.TracebackType | None,
+            ) -> None:
+        self.release()
+
+    async def __aenter__(self) -> typing.Self:
+        await self.acquire_async()
+        return self
+
+    async def __aexit__(
+            self,
+            exc_type: type | None,
+            exc_value: BaseException | None,
+            exc_traceback: types.TracebackType | None,
+            ) -> None:
+        self.release()
+
+
+def safe_read(sock: protocols.SocketLike, buffer: bytearray, size: int) -> int:
+    """Read socket into buffer and return bytes read."""
+    try:
+        data = sock.recv(size)
+
+    except (TimeoutError, BlockingIOError):
+        return 0
+
+    buffer.extend(data)
+    return len(data)
+
+
+def safe_write(sock: protocols.SocketLike, data: collections.abc.Buffer) -> int:
+    """Write buffer data into socket and return bytes written, if possible."""
+    try:
+        return sock.send(data)
+
+    except BlockingIOError as e:
+        return getattr(e, 'characters_written', 0)  # handle buffered-io-based socket-likes
+
+    except TimeoutError:
+        return 0
+
+
+def safe_timeout(sock: protocols.SocketLike, timeout: float | None) -> bool:
+    """Set timeout and return True, or return False on ValueError."""
+    try:
+        sock.settimeout(timeout)
+    except ValueError:
+        return False
+    return True
+
+
+def capture(
+        exceptions: type[Exception] | tuple[type[Exception], ...],
+        func: collections.abc.Callable, *args, **kwargs,
+        ) -> Exception | None:
+    """Call function and return matching exceptions on raise."""
+    try:
+        func(*args, **kwargs)
+    except exceptions as e:
+        return e
+    return None
+
+
+def reader(
+        sock: protocols.SocketLike,
+        buffer: bytearray,
+        size: RecvSize,
+        ) -> collections.abc.Generator[None, None, None]:
+    """Fully consume bytes from the socket."""
+    read = safe_read
+    sizes = (size,) if isinstance(size, int) else size(buffer)
+    for sz in sizes:
+        if sz:
+            while sz := (sz - read(sock, buffer, sz)):
+                yield
+
+
+def writer(sock: protocols.SocketLike, data: SendPayload) -> collections.abc.Generator[None, None, None]:
+    """Fully write iterable of buffer into socket as an iterator."""
+    write = safe_write
+    chunks = (data,) if isinstance(data, protocols.SizedBuffer) else data
+    for chunk in chunks:
+        if chunk:
+            offset = 0
+            while not (offset := write(sock, chunk)):  # write first chunk without slicing
+                yield
+
+            if offset >= len(chunk):  # finished
+                continue
+
+            yield
+            buffer = memoryview(chunk)[offset:]  # turn slicing into zero-copy
+            while not (offset := write(sock, buffer)) or (buffer := buffer[offset:]):
+                yield
+
+
+@contextlib.contextmanager
+def lock_timeout(lock: CrossLock, timeout: float | None = None) -> collections.abc.Generator[float | None, None, None]:
+    """Acquire lock and calculate deadline for timeout."""
+    timeout = socket.getdefaulttimeout() if timeout is None else timeout
+    deadline = time.monotonic() + timeout if timeout else None
+    if not lock.acquire(timeout=timeout):
+        msg = 'timeout waiting for a concurrent operation'
+        raise TimeoutError(msg)
+
+    try:
+        yield deadline
+
+    finally:
+        lock.release()
+
+
+def timeout_checker(deadline: float, error: str | None) -> collections.abc.Callable[[], float]:
+    """Get function returning seconds left before monotonic timestamp, optionally raising on timeout."""
+
+    def func() -> float:
+        if (seconds := deadline - time.monotonic()) > 0:
+            return seconds
+
+        if error:
+            raise TimeoutError(error)
+
+        return 0
+
+    return func

socketwrapper/_varint.py

@@ -0,0 +1,94 @@
+r"""SQLite4 varint tools.
+
+Implementation of SQLite4 Variable-Length Integers
+==================================================
+
+Based on spec at https://sqlite.org/src4/doc/trunk/www/varint.wiki
+
+Examples
+--------
+>>> tests = {
+...     b'\x00': 0,
+...     b'\xf0': 240,
+...     b'\xf8\xff': 2287,
+...     b'\xf9\xff\xff': 67823,
+...     b'\xfa\xff\xff\xff': 2**24 - 1,
+...     b'\xfb\xff\xff\xff\xff': 2**32 - 1,
+...     b'\xfc\xff\xff\xff\xff\xff': 2**40 - 1,
+...     b'\xfd\xff\xff\xff\xff\xff\xff': 2**48 - 1,
+...     b'\xfe\xff\xff\xff\xff\xff\xff\xff': 2**56 - 1,
+...     b'\xff\xff\xff\xff\xff\xff\xff\xff\xff': 2**64 - 1,
+...     }
+
+>>> [peek(i[:1]) for i in tests] == [len(i) for i in tests]
+True
+
+>>> [dumps(i) for i in tests.values()] == list(tests)
+True
+
+>>> [loads(i) for i in tests] == list(tests.values())
+True
+
+>>> dumps(2**64)
+Traceback (most recent call last):
+...
+OverflowError: int too big to convert
+
+"""
+
+import collections.abc
+import typing
+
+
+class VarIntBytes(typing.Protocol):
+    @typing.overload
+    def __getitem__(self, key: typing.SupportsIndex, /) -> int: ...
+
+    @typing.overload
+    def __getitem__(self, key: slice, /) -> (
+        collections.abc.Iterable[typing.SupportsIndex]
+        | typing.SupportsBytes
+        | collections.abc.Buffer
+        ): ...
+
+
+def peek(data: VarIntBytes | int) -> int:
+    """Get total varint bytesize from its first byte."""
+    start = data if isinstance(data, int) else data[0]
+    return (
+        1 if start < 0xF1 else
+        2 if start < 0xF9 else
+        3 if start < 0xFA else
+        start - 0xF6
+        )
+
+
+def dumps(value: int) -> bytes:
+    """Encode given unsigned integer as varint bytes."""
+    return (
+        value.to_bytes() if value < 0xF1 else
+        (0xF010 + value).to_bytes(2) if value < 0x8F0 else
+        (0xF8F710 + value).to_bytes(3) if value < 0x108F0 else
+        (0xFA000000 | value).to_bytes(4) if value < 0x1000000 else
+        (0xFB00000000 | value).to_bytes(5) if value < 0x100000000 else
+        (0xFC0000000000 | value).to_bytes(6) if value < 0x10000000000 else
+        (0xFD000000000000 | value).to_bytes(7) if value < 0x1000000000000 else
+        (0xFE00000000000000 | value).to_bytes(8) if value < 0x100000000000000 else
+        (0xFF0000000000000000 + value).to_bytes(9)
+        )
+
+
+def loads(data: VarIntBytes) -> int:
+    """Decode given varint bytes as integer."""
+    start = data[0]
+    return (
+        start if start < 0xF1 else
+        int.from_bytes(data[:2]) - 0xF010 if start < 0xF9 else
+        int.from_bytes(data[:3]) - 0xF8F710 if start < 0xFA else
+        int.from_bytes(data[1:start - 0xF6])
+        )
+
+
+if __name__ == '__main__':
+    import doctest
+    doctest.testmod(verbose=True, raise_on_error=True, exclude_empty=True, optionflags=doctest.ELLIPSIS)

socketwrapper/framing.py

@@ -0,0 +1,143 @@
+"""Standard message framing implementations for socketwrapper."""
+
+import abc
+import collections.abc
+import io
+import multiprocessing.connection as mp_conn
+import struct
+
+from . import _varint, protocols
+
+
+class MessageFramingBase[R, W](abc.ABC):
+    """Base abstract class for message framing implementations."""
+
+    @classmethod
+    @abc.abstractmethod
+    def frames(cls, buffer: bytearray) -> collections.abc.Iterable[int]:
+        """Get iterator with required read sizes."""
+
+    @classmethod
+    @abc.abstractmethod
+    def loads(cls, buffer: bytearray) -> R:
+        """Get frame payload from full buffer."""
+
+    @classmethod
+    @abc.abstractmethod
+    def dumps(cls, data: W) -> collections.abc.Iterable[collections.abc.Buffer]:
+        """Get iterable with frame parts (usually header and payload)."""
+
+
+class VarIntBytes(MessageFramingBase):
+    """VarInt-headered bytes message frames.
+
+    This is the optimum possible framing for bytes, using SQLite4
+    variable-length integers for the size header:
+    - Header size is variable length from 1 to 9 bytes.
+    - Maximum message length of 2**64 - 1.
+    - 2 or 3 reads in total.
+
+    """
+
+    @classmethod
+    def frames(cls, buffer: bytearray) -> collections.abc.Generator[int, None, None]:
+        """Iterate read size requests based on buffer contents."""
+        yield 1
+        yield _varint.peek(buffer) - 1
+        yield _varint.loads(buffer)
+
+    @classmethod
+    def loads(cls, buffer: bytearray) -> bytes:
+        """Get frame payload from full buffer."""
+        return bytes(buffer[_varint.peek(buffer):])
+
+    @classmethod
+    def dumps[T: protocols.SizedBuffer](cls, data: T) -> tuple[bytes, T]:
+        """Get tuple with header and payload."""
+        return _varint.dumps(len(data)), data
+
+
+class MultiprocessingFramingBase(MessageFramingBase):
+    """Base partially-abstract class for multiprocessing connection message frames.
+
+    This framing implementation uses multiprocessing connection bytes messages:
+    - Header of 4 or 12 bytes.
+    - Maximum message length of 2**64 - 1.
+    - 2 or 3 reads in total.
+
+    """
+
+    class ReadableConnection(mp_conn.Connection):
+        """Multiprocessing connection with buffered IO."""
+
+        def __init__(self, buffer: collections.abc.Buffer = b'') -> None:
+            """Initialize as readable with data."""
+            super().__init__(0, readable=True)
+            self.buffer = memoryview(buffer)
+
+        def _recv(self, size: int) -> io.BytesIO:
+            """Read data from buffer."""
+            data, self.buffer = self.buffer[:size], self.buffer[size:]
+            return io.BytesIO(data)
+
+        def _close(self) -> None:
+            """Flag as closed."""
+            self._handle = None
+
+    class WritableConnection(mp_conn.Connection):
+        """Multiprocessing connection with buffered IO."""
+
+        def __init__(self) -> None:
+            """Initialize as writable."""
+            super().__init__(0, writable=True)
+            self.chunks = []
+
+        def _send(self, buf: bytes) -> None:
+            """Record data chunk."""
+            self.chunks.append(buf)
+
+        def _close(self) -> None:
+            """Flag as closed."""
+            self._handle = None
+
+    @classmethod
+    def frames(cls, buffer: bytearray) -> collections.abc.Generator[int, None, None]:
+        """Iterate read size requests based on buffer contents."""
+        yield 4
+        size, = struct.unpack('!i', buffer)
+        if size == -1:
+            yield 8
+            size, = struct.unpack('!Q', buffer[4:])
+        yield size
+
+
+class MultiprocessingBytes(MultiprocessingFramingBase):
+    """Multiprocessing connection bytes message frames."""
+
+    @classmethod
+    def loads(cls, buffer: bytearray) -> bytes:
+        """Get frame payload from full buffer."""
+        return cls.ReadableConnection(buffer).recv_bytes()
+
+    @classmethod
+    def dumps(cls, data: collections.abc.Buffer) -> list[bytes]:
+        """Get tuple with message data."""
+        c = cls.WritableConnection()
+        c.send_bytes(data)
+        return c.chunks
+
+
+class Multiprocessing[T](MultiprocessingFramingBase):
+    """Multiprocessing connection pickled message frames."""
+
+    @classmethod
+    def loads(cls, buffer: bytearray) -> T:
+        """Get unserialized frame payload from full buffer."""
+        return cls.ReadableConnection(buffer).recv()
+
+    @classmethod
+    def dumps(cls, data: T) -> list[bytes]:
+        """Get tuple with serialized message data."""
+        c = cls.WritableConnection()
+        c.send(data)
+        return c.chunks

socketwrapper/protocols.py

@@ -0,0 +1,39 @@
+"""Standard python protocols used by socketwrapper."""
+
+import collections.abc
+import typing
+
+
+@typing.runtime_checkable
+class SizedBuffer(typing.Protocol):
+    """Protocol for buffers with size."""
+
+    def __buffer__(self, flags: int, /) -> memoryview: ...
+    def __len__(self) -> int: ...
+
+
+@typing.runtime_checkable
+class SocketLike(typing.Protocol):
+    """Protocol for socket-like objects accepted by socketwrapper socket classes."""
+
+    def fileno(self) -> int: ...
+    def send(self, data: collections.abc.Buffer, /) -> int: ...
+    def recv(self, bufsize: int, /) -> bytes: ...
+    def settimeout(self, timeout: float | None, /) -> None: ...
+    def close(self) -> None: ...
+
+
+@typing.runtime_checkable
+class HasFileno(typing.Protocol):
+    """Protocol for objects backed by a file descriptor."""
+
+    def fileno(self) -> int: ...
+
+
+@typing.runtime_checkable
+class MessageFraming[R, W](typing.Protocol):
+    """Protocol for message framing implementations accepted by socketwrapper message classes."""
+
+    def frames(self, buffer: bytearray, /) -> collections.abc.Iterable[int]: ...
+    def loads(self, buffer: bytearray, /) -> R: ...
+    def dumps(self, data: W, /) -> collections.abc.Iterable[SizedBuffer]: ...

socketwrapper-0.0.1.dist-info/METADATA

@@ -0,0 +1,294 @@
+Metadata-Version: 2.4
+Name: socketwrapper
+Version: 0.0.1
+Summary: high level socket and pipe wrapper
+Author: Felipe A Hernandez
+Author-email: ergoithz@gmail.com
+License: MIT License
+        
+        Copyright (c) 2025 Felipe A Hernandez <ergoithz@gmail.com>
+        
+        Permission is hereby granted, free of charge, to any person obtaining a copy
+        of this software and associated documentation files (the "Software"), to deal
+        in the Software without restriction, including without limitation the rights
+        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+        copies of the Software, and to permit persons to whom the Software is
+        furnished to do so, subject to the following conditions:
+        
+        The above copyright notice and this permission notice shall be included in all
+        copies or substantial portions of the Software.
+        
+        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+        SOFTWARE.
+        
+Project-URL: homepage, https://gitlab.com/ergoithz/socketwrapper
+Project-URL: issue-tracker, https://gitlab.com/ergoithz/socketwrapper/-/issues
+Project-URL: release-notes, https://gitlab.com/ergoithz/socketwrapper/-/releases
+Project-URL: issue-new, https://gitlab.com/ergoithz/socketwrapper/-/issues/new
+Project-URL: donations, https://ko-fi.com/s26me
+Keywords: socket,pipe,ipc,asyncio
+Classifier: Framework :: AsyncIO
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Software Development :: Libraries
+Classifier: Topic :: System :: Networking
+Requires-Python: >=3.12
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: typing-extensions>=4.15; python_version < "3.13"
+Provides-Extra: dev
+Requires-Dist: coverage; extra == "dev"
+Requires-Dist: ruff; extra == "dev"
+Requires-Dist: wheel; extra == "dev"
+Requires-Dist: yapf; extra == "dev"
+Dynamic: license-file
+
+# socketwrapper
+
+This package provides high level wrappers for sockets and pipes:
+- Thread-safe within both and between threads and asyncio realms.
+- Managed sync recv and send operations with timeouts.
+- Native asyncio recv_async and send_async operations.
+- Pluggable message protocol (headered variable length data) supporting
+  header parsing, serialization and deserialization.
+- Pluggable I/O protocol (file descriptor still required due asyncio).
+
+## Motivation
+
+There aren't a ton of high level asyncio socket wrappers out there providing
+all we need for IPC: header/payload logic and support for both sockets and pipes.
+
+Most implementations either don't have pluggable messaging protocols or it
+doesn't support variable-size headers.
+
+No other implementation is thread-safe across both asyncio and threading realms.
+
+No other implementation directly supports wrapping multiprocessing Connections.
+
+## Installation
+
+```sh
+uv pip install socketwrapper
+```
+
+## Documentation
+
+None other than this README, life's too short and I'm too busy with real life,
+if you need better documentation consider donating to
+[my ko-fi](https://ko-fi.com/s26me) stating that as the tip message,
+check out how my docs look like at [mstache docs](https://mstache.readthedocs.io/en/latest/)
+and [uactor docs](https://mstache.readthedocs.io/en/latest/).
+
+### Puggable I/O: SocketLike protocol
+
+The `socketwrapper.SocketLike` protocol, an small subset the socket interface,
+is all what's required for any object to be wrap-able by `socketwrapper`.
+
+```py
+@typing.runtime_checkable
+class SocketLike(typing.Protocol):
+    """Protocol for socket-like objects accepted by socketwrapper socket classes."""
+
+    def fileno(self) -> int: ...
+    def send(self, data: collections.abc.Buffer, /) -> int: ...
+    def recv(self, bufsize: int, /) -> bytes: ...
+    def settimeout(self, timeout: float | None, /) -> None: ...
+    def close(self) -> None: ...
+```
+
+Special attention to:
+- [fileno](https://docs.python.org/3.14/library/socket.html#socket.socket.fileno)
+  has to be a valid OS file descriptor.
+- [settimeout](https://docs.python.org/3.14/library/socket.html#socket.socket.settimeout)
+  only requires support for `settimeout(.0)` (
+  [non-blocking semantics](https://docs.python.org/3.14/library/socket.html#notes-on-socket-timeouts)
+  ), raising [ValueError](https://docs.python.org/3.14/library/exceptions.html#ValueError)
+  for any other value is fully supported in which case
+  [selectors.DefaultSelector](https://docs.python.org/3.14/library/selectors.html#selectors.DefaultSelector)
+  will be used for synchronous operations.
+
+## Usage
+
+### Simple IPC with pipe
+
+```python
+import os
+import socketwrapper
+
+with socketwrapper.pipe(framing=True) as (parent_reader, child_writer):
+    child_writer.inheritable = True
+    child_pid = os.fork()  # replace with your own process fork/spawn logic
+    child_writer.inheritable = False  # important, prevent socket leaks!
+
+    if child_pid:
+        print(f'Message {parent_reader.recv()!r} received')
+    else:
+        child_writer.send(b'Hello world!')
+```
+
+```
+Message b'Hello world!' received
+```
+
+### Simple asyncio IPC with pipe
+
+```python
+import asyncio
+import os
+import socketwrapper
+
+async def parent(readable: socketwrapper.MessageReader) -> None:
+    print(f'Message {await readable.recv_async()!r} received')
+
+async def child(writable: socketwrapper.MessageWriter) -> None:
+    await writable.send_async(b'Hello world!')
+
+with socketwrapper.pipe(framing=True) as (parent_reader, child_writer):
+    child_writer.inheritable = True
+    child_pid = os.fork()  # replace with your own process fork/spawn logic
+    child_writer.inheritable = False  # important, prevent socket leaks!
+
+    asyncio.run(parent(parent_reader) if child_pid else child(child_writer))
+```
+
+```
+Message b'Hello world!' received
+```
+
+### Simple bidirectional IPC with socketpair
+
+```python
+import os
+import socketwrapper
+
+with socketwrapper.socketpair(framing=True) as (parent_duplex, child_duplex):
+    child_duplex.inheritable = True
+    child_pid = os.fork()  # replace with your own process fork/spawn logic
+    child_duplex.inheritable = False  # important, prevent socket leaks!
+
+    if child_pid:
+        parent_duplex.send(b'Hello child!')
+        print(f'Message {parent_duplex.recv()!r} received in parent')
+
+    else:
+        print(f'Message {child_duplex.recv()!r} received in child')
+        child_duplex.send(b'Hello parent!')
+```
+
+```
+Message b'Hello child!' received in child
+Message b'Hello parent!' received in parent
+```
+
+### Socketwrapper with multiprocessing.Pipe and asyncio
+
+```py
+import asyncio
+import multiprocessing
+import multiprocessing.connection
+import socketwrapper
+import socketwrapper.framing
+
+def child(conn: multiprocessing.connection.Connection) -> None:
+
+    async def main() -> None:
+        with socketwrapper.MessageDuplex(conn, framing=socketwrapper.framing.MultiprocessingBytes) as child_duplex:
+            print(f'Message {await child_duplex.recv_async()!r} received in child')
+            await child_duplex.send_async(b'Hello parent!')
+
+    asyncio.run(main())
+
+parent_conn, child_conn = multiprocessing.Pipe()
+with parent_conn, child_conn:
+    child_process = multiprocessing.Process(target=child, args=(child_conn,))
+    child_process.start()
+
+    parent_conn.send_bytes(b'Hello child!')
+    print(f'Message {parent_conn.recv_bytes()!r} received in parent')
+    child_process.join(1)
+```
+
+```
+Message b'Hello child!' received in child
+Message b'Hello parent!' received in parent
+```
+
+### Custom socketwrapper framing with progress
+
+```py
+import collections.abc
+import itertools
+import os
+import socketwrapper
+import socketwrapper.framing
+
+def progress(arrow: str, size: int, min_chunk: int = 1024) -> collections.abc.Generator[int, None, None]:
+    part_size = max(min_chunk, size // 100)
+    full_parts, last_size = divmod(size, part_size)
+    percent = 100 / (full_parts + 1 if last_size else full_parts)
+
+    for i in range(full_parts):
+        print(f'{arrow} {i * percent:6.2f}%')
+        yield part_size
+
+    if last_size:
+        print(f'{arrow} {full_parts * percent:6.2f}%')
+        yield last_size
+
+    print(f'{arrow} 100%')
+
+class ProgressFraming(socketwrapper.framing.VarIntBytes):
+
+    @classmethod
+    def frames(cls, buffer: bytearray) -> collections.abc.Generator[int, None, None]:
+        frames = super().frames(buffer)
+        yield from itertools.islice(frames, 2)
+        yield from progress('>', next(frames))
+
+    @classmethod
+    def dumps(cls, data: bytes) -> collections.abc.Generator[memoryview, None, None]:
+        buffer = memoryview(b''.join(super().dumps(data)))
+        for size in progress('<', len(buffer)):
+            chunk, buffer = buffer[:size], buffer[size:]
+            yield chunk
+
+with socketwrapper.socketpair(framing=ProgressFraming) as (parent_duplex, child_duplex):
+    child_duplex.inheritable = True
+    child_pid = os.fork()  # replace with your own multiprocessing fork logic
+    child_duplex.inheritable = False  # important, prevent socket leaks!
+
+    if child_pid:
+        payload = os.urandom(1024) * 1024
+        print(f'Sending {len(payload)} bytes!')
+        parent_duplex.send(payload)
+    else:
+        print(f'Received {len(child_duplex.recv())} bytes!')
+```
+```
+Sending 1048576 bytes!
+<   0.00%
+<   0.99%
+<   1.98%
+<   2.97%
+...
+>   0.99%
+<  13.86%
+>   1.98%
+<  14.85%
+...
+>  91.09%
+<  99.01%
+>  92.08%
+< 100%
+...
+>  98.02%
+>  99.01%
+> 100%
+Received 1048576 bytes!
+```

socketwrapper-0.0.1.dist-info/RECORD

@@ -0,0 +1,12 @@
+socketwrapper/__init__.py,sha256=XZxrB-6eF3Ntn0DDsgE0gRM0cvE7i1aNMo3X_SnzDAg,15463
+socketwrapper/_utils.py,sha256=1iTOnnII0McX4p_ikjQ8INdz-DDigSR6hMrQ2He7eKw,8568
+socketwrapper/_varint.py,sha256=VPDqH8vS4Oixutk8ZWsmqBMDajjODrZ-7xM9Vm3-u1E,2786
+socketwrapper/framing.py,sha256=YL8QHR9LbuCOC4xJ1wH0mrK5pFmiJn7jec8U0PDjtCU,4563
+socketwrapper/protocols.py,sha256=xpGNqxRXiSy0LiGZ9WamZsJy2w6wmvv-Kwy9oCHyTyQ,1225
+socketwrapper/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+socketwrapper-0.0.1.dist-info/licenses/LICENSE,sha256=tIR671OctNcul-PHVKPNff7pMLcvO5z1AFeDNH9nb-0,1096
+socketwrapper-0.0.1.dist-info/METADATA,sha256=vJoJobF2fogbgqUQb1zyjDiaA50zijKKWtl_AtYWiDs,9932
+socketwrapper-0.0.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+socketwrapper-0.0.1.dist-info/top_level.txt,sha256=-NM1Y47ygwdVOVqMF16S23zt2P1FluEV-3z4Jzmavvo,14
+socketwrapper-0.0.1.dist-info/zip-safe,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
+socketwrapper-0.0.1.dist-info/RECORD,,

socketwrapper-0.0.1.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (80.9.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

socketwrapper-0.0.1.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Felipe A Hernandez <ergoithz@gmail.com>
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

socketwrapper-0.0.1.dist-info/top_level.txt

@@ -0,0 +1 @@
+socketwrapper

socketwrapper-0.0.1.dist-info/zip-safe

@@ -0,0 +1 @@
+