91 lines
3.1 KiB
Python
91 lines
3.1 KiB
Python
"""Binary TCP protocol shared by the K209 remote server and Python client."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import socket
|
|
import struct
|
|
from typing import BinaryIO
|
|
|
|
import numpy as np
|
|
|
|
COMMAND_IDENTITY = b"I"
|
|
COMMAND_LIMITS = b"L"
|
|
COMMAND_CONFIGURE = b"C"
|
|
COMMAND_ACQUIRE = b"A"
|
|
|
|
STATUS_OK = b"O"
|
|
STATUS_ERROR = b"E"
|
|
|
|
DEFAULT_REMOTE_HOST = "127.0.0.1"
|
|
DEFAULT_REMOTE_PORT = 50209
|
|
|
|
CONFIG_STRUCT = struct.Struct("!ddIdd")
|
|
LIMITS_STRUCT = struct.Struct("!ddddIdd")
|
|
U32_STRUCT = struct.Struct("!I")
|
|
|
|
|
|
def recv_exact(stream: socket.socket | BinaryIO, size: int) -> bytes:
|
|
"""Read exactly `size` bytes from a socket-like object."""
|
|
chunks: list[bytes] = []
|
|
remaining = int(size)
|
|
while remaining > 0:
|
|
chunk = stream.recv(remaining) if isinstance(stream, socket.socket) else stream.read(remaining)
|
|
if not chunk:
|
|
raise ConnectionError(f"K209 remote connection closed with {remaining} bytes pending")
|
|
chunks.append(chunk)
|
|
remaining -= len(chunk)
|
|
return b"".join(chunks)
|
|
|
|
|
|
def send_all(stream: socket.socket | BinaryIO, payload: bytes) -> None:
|
|
"""Write all payload bytes to a socket-like object."""
|
|
if isinstance(stream, socket.socket):
|
|
stream.sendall(payload)
|
|
return
|
|
stream.write(payload)
|
|
|
|
|
|
def send_u32(stream: socket.socket | BinaryIO, value: int) -> None:
|
|
"""Send one network-order uint32."""
|
|
send_all(stream, U32_STRUCT.pack(int(value)))
|
|
|
|
|
|
def recv_u32(stream: socket.socket | BinaryIO) -> int:
|
|
"""Read one network-order uint32."""
|
|
return int(U32_STRUCT.unpack(recv_exact(stream, U32_STRUCT.size))[0])
|
|
|
|
|
|
def send_error(stream: socket.socket | BinaryIO, message: str) -> None:
|
|
"""Send protocol error response."""
|
|
payload = str(message).encode("utf-8", errors="replace")
|
|
send_all(stream, STATUS_ERROR)
|
|
send_u32(stream, len(payload))
|
|
send_all(stream, payload)
|
|
|
|
|
|
def read_status(stream: socket.socket | BinaryIO) -> None:
|
|
"""Read response status and raise remote error when needed."""
|
|
status = recv_exact(stream, 1)
|
|
if status == STATUS_OK:
|
|
return
|
|
if status == STATUS_ERROR:
|
|
message = recv_exact(stream, recv_u32(stream)).decode("utf-8", errors="replace")
|
|
raise RuntimeError(f"K209 remote server error: {message}")
|
|
raise RuntimeError(f"K209 remote server returned invalid status byte: {status!r}")
|
|
|
|
|
|
def send_float32_array(stream: socket.socket | BinaryIO, values: np.ndarray) -> None:
|
|
"""Send a float32 array as a length-prefixed little-endian payload."""
|
|
payload = np.asarray(values, dtype="<f4").tobytes(order="C")
|
|
send_u32(stream, len(payload))
|
|
send_all(stream, payload)
|
|
|
|
|
|
def recv_float32_array(stream: socket.socket | BinaryIO, expected_values: int) -> np.ndarray:
|
|
"""Read a length-prefixed little-endian float32 array."""
|
|
payload_size = recv_u32(stream)
|
|
expected_size = int(expected_values) * np.dtype(np.float32).itemsize
|
|
if payload_size != expected_size:
|
|
raise RuntimeError(f"K209 remote payload has {payload_size} bytes, expected {expected_size}")
|
|
return np.frombuffer(recv_exact(stream, payload_size), dtype="<f4").copy()
|