robosim.network.msgpack_utils

 1"""utils.py - Network utils"""
 2from typing import Any
 3from datetime import datetime
 4import socket
 5
 6import msgpack
 7
 8from robosim.utils import logger
 9
10_SERIALIZABLE_TYPES = (str, int, float, bool, bytes, type(None))
11
12
13def send_packet(sock: socket.socket, data: dict, close_connection: bool=False, add_timestmap: bool=True):
14    """pack the message in messagepack format + length-prefixed for TCP stream then send it through the socket"""
15    assert isinstance(data, dict), (data, type(data))
16    if add_timestmap:
17        assert "timestamp" not in data.keys(), data
18        data["timestamp"] = datetime.now().isoformat()
19    try:
20        packed_data = msgpack.packb(data)
21    except TypeError as e:
22        _find_bad(data)
23        raise e
24    packet = len(packed_data).to_bytes(4, "big") + packed_data # <SIZE-4b><packed_data>
25    logger.trace(f"Sending {len(packet)} bytes through the wire")
26    sock.sendall(packet)
27
28    if close_connection:
29        sock.close()
30
31def recv_packet(sock: socket.socket) -> dict:
32    """reads a packet from the server. First 4 bytes for len, then the rest of the message"""
33    packet_len = int.from_bytes(_recvall(sock, 4), "big")
34    packet: dict = msgpack.unpackb(_recvall(sock, packet_len), raw=False)
35    assert isinstance(packet, dict), f"packet is not dictionary: {type(packet)}"
36    return packet
37
38def _find_bad(d: Any, path: str="root"):
39    """ran in case of a TypeError in send_packet, so we know which part of the dict is not serializable"""
40    if isinstance(d, dict):
41        for k, v in d.items():
42            _find_bad(v, f"{path}[{k!r}]") # !r calls repr()
43    elif isinstance(d, (list, tuple)):
44        for i, v in enumerate(d):
45            _find_bad(v, f"{path}[{i}]")
46    elif not isinstance(d, _SERIALIZABLE_TYPES):
47        logger.error(f"Unserializable at {path}: {type(d).__name__}")
48    # the 'else' is implied because only 'good' types should be left here
49
50def _recvall(sock: socket.socket, n: int) -> bytes:
51    """recvall receives all n bytes at once from a tcp socket"""
52    buf = bytearray()
53    while len(buf) < n:
54        if not (chunk := sock.recv(n - len(buf))):
55            raise ConnectionError
56        buf.extend(chunk)
57    return bytes(buf)