rocket-welder-sdk 1.1.32__py3-none-any.whl → 1.1.33__py3-none-any.whl

rocket_welder_sdk/transport/nng_transport.py

@@ -0,0 +1,197 @@
+"""NNG transport using pynng library.
+
+NNG (nanomsg next generation) provides high-performance, scalable messaging patterns.
+Supported patterns:
+- Pub/Sub: One publisher to many subscribers
+- Push/Pull: Load-balanced distribution to workers
+"""
+
+from typing import Any, Optional, cast
+
+import pynng
+
+from .frame_sink import IFrameSink
+from .frame_source import IFrameSource
+
+
+class NngFrameSink(IFrameSink):
+    """
+    Frame sink that publishes to NNG Pub/Sub or Push/Pull pattern.
+
+    Each frame is sent as a single NNG message (no framing needed - NNG handles message boundaries).
+    """
+
+    def __init__(self, socket: Any, leave_open: bool = False):
+        """
+        Create an NNG frame sink from a socket.
+
+        Args:
+            socket: pynng socket (Publisher or Pusher)
+            leave_open: If True, doesn't close socket on close
+        """
+        self._socket: Any = socket
+        self._leave_open = leave_open
+        self._closed = False
+
+    @classmethod
+    def create_publisher(cls, url: str) -> "NngFrameSink":
+        """
+        Create an NNG Publisher frame sink bound to the specified URL.
+
+        Args:
+            url: NNG URL (e.g., "tcp://127.0.0.1:5555", "ipc:///tmp/mysocket")
+
+        Returns:
+            Frame sink ready to publish messages
+        """
+        socket = pynng.Pub0()
+        socket.listen(url)
+        return cls(socket, leave_open=False)
+
+    @classmethod
+    def create_pusher(cls, url: str, bind_mode: bool = True) -> "NngFrameSink":
+        """
+        Create an NNG Pusher frame sink.
+
+        Args:
+            url: NNG URL (e.g., "tcp://127.0.0.1:5555", "ipc:///tmp/mysocket")
+            bind_mode: If True, listens (bind); if False, dials (connect)
+
+        Returns:
+            Frame sink ready to push messages
+        """
+        socket = pynng.Push0()
+        if bind_mode:
+            socket.listen(url)
+        else:
+            socket.dial(url)
+        return cls(socket, leave_open=False)
+
+    def write_frame(self, frame_data: bytes) -> None:
+        """Write frame to NNG socket (no length prefix - NNG handles message boundaries)."""
+        if self._closed:
+            raise ValueError("Cannot write to closed sink")
+
+        self._socket.send(frame_data)
+
+    async def write_frame_async(self, frame_data: bytes) -> None:
+        """Write frame asynchronously."""
+        if self._closed:
+            raise ValueError("Cannot write to closed sink")
+
+        await self._socket.asend(frame_data)
+
+    def flush(self) -> None:
+        """Flush is a no-op for NNG (data sent immediately)."""
+        pass
+
+    async def flush_async(self) -> None:
+        """Flush asynchronously is a no-op for NNG."""
+        pass
+
+    def close(self) -> None:
+        """Close the NNG sink."""
+        if self._closed:
+            return
+        self._closed = True
+        if not self._leave_open:
+            self._socket.close()
+
+    async def close_async(self) -> None:
+        """Close the NNG sink asynchronously."""
+        self.close()
+
+
+class NngFrameSource(IFrameSource):
+    """
+    Frame source that subscribes to NNG Pub/Sub or Pull pattern.
+
+    Each NNG message is treated as a complete frame (no framing needed - NNG handles message boundaries).
+    """
+
+    def __init__(self, socket: Any, leave_open: bool = False):
+        """
+        Create an NNG frame source from a socket.
+
+        Args:
+            socket: pynng socket (Subscriber or Puller)
+            leave_open: If True, doesn't close socket on close
+        """
+        self._socket: Any = socket
+        self._leave_open = leave_open
+        self._closed = False
+
+    @classmethod
+    def create_subscriber(cls, url: str, topic: bytes = b"") -> "NngFrameSource":
+        """
+        Create an NNG Subscriber frame source connected to the specified URL.
+
+        Args:
+            url: NNG URL (e.g., "tcp://127.0.0.1:5555", "ipc:///tmp/mysocket")
+            topic: Optional topic filter (empty for all messages)
+
+        Returns:
+            Frame source ready to receive messages
+        """
+        socket = pynng.Sub0()
+        socket.subscribe(topic)
+        socket.dial(url)
+        return cls(socket, leave_open=False)
+
+    @classmethod
+    def create_puller(cls, url: str, bind_mode: bool = True) -> "NngFrameSource":
+        """
+        Create an NNG Puller frame source.
+
+        Args:
+            url: NNG URL (e.g., "tcp://127.0.0.1:5555", "ipc:///tmp/mysocket")
+            bind_mode: If True, listens (bind); if False, dials (connect)
+
+        Returns:
+            Frame source ready to pull messages
+        """
+        socket = pynng.Pull0()
+        if bind_mode:
+            socket.listen(url)
+        else:
+            socket.dial(url)
+        return cls(socket, leave_open=False)
+
+    @property
+    def has_more_frames(self) -> bool:
+        """Check if more frames available (NNG blocks waiting for messages)."""
+        return not self._closed
+
+    def read_frame(self) -> Optional[bytes]:
+        """Read frame from NNG socket (blocking)."""
+        if self._closed:
+            return None
+
+        try:
+            return cast("bytes", self._socket.recv())
+        except pynng.Closed:
+            self._closed = True
+            return None
+
+    async def read_frame_async(self) -> Optional[bytes]:
+        """Read frame asynchronously."""
+        if self._closed:
+            return None
+
+        try:
+            return cast("bytes", await self._socket.arecv())
+        except pynng.Closed:
+            self._closed = True
+            return None
+
+    def close(self) -> None:
+        """Close the NNG source."""
+        if self._closed:
+            return
+        self._closed = True
+        if not self._leave_open:
+            self._socket.close()
+
+    async def close_async(self) -> None:
+        """Close the NNG source asynchronously."""
+        self.close()

rocket_welder_sdk/transport/stream_transport.py

@@ -0,0 +1,193 @@
+"""Stream-based transport (file, memory, etc.)."""
+
+from typing import BinaryIO, Optional
+
+from .frame_sink import IFrameSink
+from .frame_source import IFrameSource
+
+
+def _write_varint(stream: BinaryIO, value: int) -> None:
+    """Write unsigned integer as varint (Protocol Buffers format)."""
+    if value < 0:
+        raise ValueError(f"Varint requires non-negative value, got {value}")
+
+    while value >= 0x80:
+        stream.write(bytes([value & 0x7F | 0x80]))
+        value >>= 7
+    stream.write(bytes([value & 0x7F]))
+
+
+def _read_varint(stream: BinaryIO) -> int:
+    """Read varint from stream and decode to unsigned integer."""
+    result = 0
+    shift = 0
+
+    while True:
+        if shift >= 35:  # Max 5 bytes for uint32
+            raise ValueError("Varint too long (corrupted stream)")
+
+        byte_data = stream.read(1)
+        if not byte_data:
+            raise EOFError("Unexpected end of stream reading varint")
+
+        byte = byte_data[0]
+        result |= (byte & 0x7F) << shift
+        shift += 7
+
+        if not (byte & 0x80):
+            break
+
+    return result
+
+
+class StreamFrameSink(IFrameSink):
+    """
+    Frame sink that writes to a BinaryIO stream (file, memory, etc.).
+
+    Each frame is prefixed with its length (varint encoding) for frame boundary detection.
+    Format: [varint length][frame data]
+    """
+
+    def __init__(self, stream: BinaryIO, leave_open: bool = False):
+        """
+        Create a stream-based frame sink.
+
+        Args:
+            stream: Binary stream to write to
+            leave_open: If True, doesn't close stream on close
+        """
+        self._stream = stream
+        self._leave_open = leave_open
+        self._closed = False
+
+    def write_frame(self, frame_data: bytes) -> None:
+        """Write frame data to stream with varint length prefix."""
+        if self._closed:
+            raise ValueError("Cannot write to closed sink")
+
+        # Write frame length as varint
+        _write_varint(self._stream, len(frame_data))
+
+        # Write frame data
+        self._stream.write(frame_data)
+
+    async def write_frame_async(self, frame_data: bytes) -> None:
+        """Write frame data to stream asynchronously."""
+        # For regular streams, just use synchronous write
+        # If stream supports async, could use aiofiles
+        self.write_frame(frame_data)
+
+    def flush(self) -> None:
+        """Flush buffered data to stream."""
+        if not self._closed:
+            self._stream.flush()
+
+    async def flush_async(self) -> None:
+        """Flush buffered data to stream asynchronously."""
+        self.flush()
+
+    def close(self) -> None:
+        """Close the sink."""
+        if self._closed:
+            return
+        self._closed = True
+        if not self._leave_open:
+            self._stream.close()
+
+    async def close_async(self) -> None:
+        """Close the sink asynchronously."""
+        self.close()
+
+
+class StreamFrameSource(IFrameSource):
+    """
+    Frame source that reads from a BinaryIO stream (file, memory, etc.).
+
+    Reads frames prefixed with varint length for frame boundary detection.
+    Format: [varint length][frame data]
+    """
+
+    def __init__(self, stream: BinaryIO, leave_open: bool = False):
+        """
+        Create a stream-based frame source.
+
+        Args:
+            stream: Binary stream to read from
+            leave_open: If True, doesn't close stream on close
+        """
+        self._stream = stream
+        self._leave_open = leave_open
+        self._closed = False
+
+    @property
+    def has_more_frames(self) -> bool:
+        """Check if more data available in stream."""
+        if self._closed:
+            return False
+        current_pos = self._stream.tell()
+        # Try seeking to end to check size
+        try:
+            self._stream.seek(0, 2)  # Seek to end
+            end_pos = self._stream.tell()
+            self._stream.seek(current_pos)  # Restore position
+            return current_pos < end_pos
+        except OSError:
+            # Stream not seekable, assume data available
+            return True
+
+    def read_frame(self) -> Optional[bytes]:
+        """
+        Read frame from stream with varint length-prefix framing.
+
+        Returns:
+            Frame data bytes, or None if end of stream
+        """
+        if self._closed:
+            return None
+
+        # Check if stream has data (for seekable streams)
+        if hasattr(self._stream, "tell") and hasattr(self._stream, "seek"):
+            try:
+                current_pos = self._stream.tell()
+                self._stream.seek(0, 2)  # Seek to end
+                end_pos = self._stream.tell()
+                self._stream.seek(current_pos)  # Restore position
+                if current_pos >= end_pos:
+                    return None
+            except OSError:
+                pass  # Stream not seekable, continue
+
+        # Read frame length (varint)
+        try:
+            frame_length = _read_varint(self._stream)
+        except EOFError:
+            return None
+
+        if frame_length == 0:
+            return b""
+
+        # Read frame data
+        frame_data = self._stream.read(frame_length)
+        if len(frame_data) != frame_length:
+            raise EOFError(
+                f"Unexpected end of stream while reading frame. Expected {frame_length} bytes, got {len(frame_data)}"
+            )
+
+        return frame_data
+
+    async def read_frame_async(self) -> Optional[bytes]:
+        """Read frame from stream asynchronously."""
+        # For regular streams, just use synchronous read
+        return self.read_frame()
+
+    def close(self) -> None:
+        """Close the source."""
+        if self._closed:
+            return
+        self._closed = True
+        if not self._leave_open:
+            self._stream.close()
+
+    async def close_async(self) -> None:
+        """Close the source asynchronously."""
+        self.close()

rocket_welder_sdk/transport/tcp_transport.py

@@ -0,0 +1,154 @@
+"""TCP transport with length-prefix framing."""
+
+import contextlib
+import socket
+import struct
+from typing import Optional
+
+from .frame_sink import IFrameSink
+from .frame_source import IFrameSource
+
+
+class TcpFrameSink(IFrameSink):
+    """
+    Frame sink that writes to a TCP connection with length-prefix framing.
+
+    Each frame is prefixed with a 4-byte little-endian length header.
+
+    Frame format: [Length: 4 bytes LE][Frame Data: N bytes]
+    """
+
+    def __init__(self, sock: socket.socket, leave_open: bool = False):
+        """
+        Create a TCP frame sink.
+
+        Args:
+            sock: TCP socket to write to
+            leave_open: If True, doesn't close socket on close
+        """
+        self._socket = sock
+        self._leave_open = leave_open
+        self._closed = False
+
+    def write_frame(self, frame_data: bytes) -> None:
+        """Write frame with length prefix to TCP socket."""
+        if self._closed:
+            raise ValueError("Cannot write to closed sink")
+
+        # Write 4-byte length prefix (little-endian)
+        length_prefix = struct.pack("<I", len(frame_data))
+        self._socket.sendall(length_prefix)
+
+        # Write frame data
+        self._socket.sendall(frame_data)
+
+    async def write_frame_async(self, frame_data: bytes) -> None:
+        """Write frame asynchronously (uses sync socket for now)."""
+        self.write_frame(frame_data)
+
+    def flush(self) -> None:
+        """Flush is a no-op for TCP (data sent immediately)."""
+        pass
+
+    async def flush_async(self) -> None:
+        """Flush asynchronously is a no-op for TCP."""
+        pass
+
+    def close(self) -> None:
+        """Close the TCP sink."""
+        if self._closed:
+            return
+        self._closed = True
+        if not self._leave_open:
+            with contextlib.suppress(OSError):
+                self._socket.shutdown(socket.SHUT_WR)
+            self._socket.close()
+
+    async def close_async(self) -> None:
+        """Close the TCP sink asynchronously."""
+        self.close()
+
+
+class TcpFrameSource(IFrameSource):
+    """
+    Frame source that reads from a TCP connection with length-prefix framing.
+
+    Each frame is prefixed with a 4-byte little-endian length header.
+
+    Frame format: [Length: 4 bytes LE][Frame Data: N bytes]
+    """
+
+    def __init__(self, sock: socket.socket, leave_open: bool = False):
+        """
+        Create a TCP frame source.
+
+        Args:
+            sock: TCP socket to read from
+            leave_open: If True, doesn't close socket on close
+        """
+        self._socket = sock
+        self._leave_open = leave_open
+        self._closed = False
+        self._end_of_stream = False
+
+    @property
+    def has_more_frames(self) -> bool:
+        """Check if more frames available."""
+        return not self._closed and not self._end_of_stream
+
+    def read_frame(self) -> Optional[bytes]:
+        """Read frame with length prefix from TCP socket."""
+        if self._closed or self._end_of_stream:
+            return None
+
+        # Read 4-byte length prefix
+        length_data = self._recv_exactly(4)
+        if length_data is None or len(length_data) < 4:
+            self._end_of_stream = True
+            return None
+
+        frame_length = struct.unpack("<I", length_data)[0]
+
+        if frame_length == 0:
+            return b""
+
+        if frame_length > 100 * 1024 * 1024:  # 100 MB sanity check
+            raise ValueError(f"Frame length {frame_length} exceeds maximum")
+
+        # Read frame data
+        frame_data = self._recv_exactly(frame_length)
+        if frame_data is None or len(frame_data) < frame_length:
+            self._end_of_stream = True
+            raise ValueError(
+                f"Incomplete frame data: expected {frame_length}, got {len(frame_data) if frame_data else 0}"
+            )
+
+        return frame_data
+
+    async def read_frame_async(self) -> Optional[bytes]:
+        """Read frame asynchronously (uses sync socket for now)."""
+        return self.read_frame()
+
+    def _recv_exactly(self, n: int) -> Optional[bytes]:
+        """Receive exactly n bytes from socket."""
+        data = b""
+        while len(data) < n:
+            chunk = self._socket.recv(n - len(data))
+            if not chunk:
+                return data if data else None
+            data += chunk
+        return data
+
+    def close(self) -> None:
+        """Close the TCP source."""
+        if self._closed:
+            return
+        self._closed = True
+        if not self._leave_open:
+            with contextlib.suppress(OSError):
+                self._socket.shutdown(socket.SHUT_RD)
+            self._socket.close()
+
+    async def close_async(self) -> None:
+        """Close the TCP source asynchronously."""
+        self.close()