caudal/programa/cartera/msgpack_min.py

"""msgpack_min.py — MessagePack mínimo, lo justo para el canal de precios de BME (sin dependencias)."""
import datetime as dt
import struct


def empaquetar(o) -> bytes:
    if o is None:
        return b"\xc0"
    if o is True:
        return b"\xc3"
    if o is False:
        return b"\xc2"
    if isinstance(o, int):
        if 0 <= o < 128:
            return bytes([o])
        if -32 <= o < 0:
            return struct.pack("b", o)
        return b"\xd3" + struct.pack(">q", o)
    if isinstance(o, float):
        return b"\xcb" + struct.pack(">d", o)
    if isinstance(o, str):
        b = o.encode("utf-8")
        n = len(b)
        cab = bytes([0xa0 | n]) if n < 32 else b"\xd9" + bytes([n]) if n < 256 else b"\xda" + struct.pack(">H", n)
        return cab + b
    if isinstance(o, (list, tuple)):
        n = len(o)
        cab = bytes([0x90 | n]) if n < 16 else b"\xdc" + struct.pack(">H", n)
        return cab + b"".join(empaquetar(x) for x in o)
    if isinstance(o, dict):
        n = len(o)
        cab = bytes([0x80 | n]) if n < 16 else b"\xde" + struct.pack(">H", n)
        return cab + b"".join(empaquetar(k) + empaquetar(v) for k, v in o.items())
    raise TypeError(type(o))


def desempaquetar(b: bytes, i: int = 0):
    t = b[i]
    i += 1
    if t < 0x80:
        return t, i
    if t >= 0xe0:
        return t - 256, i
    if 0x80 <= t <= 0x8f:
        return _mapa(b, i, t & 0x0f)
    if 0x90 <= t <= 0x9f:
        return _lista(b, i, t & 0x0f)
    if 0xa0 <= t <= 0xbf:
        n = t & 0x1f
        return b[i:i + n].decode("utf-8"), i + n
    fijos = {0xc0: None, 0xc2: False, 0xc3: True}
    if t in fijos:
        return fijos[t], i
    tam = {0xcc: ">B", 0xcd: ">H", 0xce: ">I", 0xcf: ">Q", 0xd0: ">b", 0xd1: ">h", 0xd2: ">i", 0xd3: ">q", 0xca: ">f", 0xcb: ">d"}
    if t in tam:
        f = tam[t]
        n = struct.calcsize(f)
        return struct.unpack(f, b[i:i + n])[0], i + n
    if t in (0xd9, 0xda, 0xdb, 0xc4, 0xc5, 0xc6):
        f = {0xd9: ">B", 0xda: ">H", 0xdb: ">I", 0xc4: ">B", 0xc5: ">H", 0xc6: ">I"}[t]
        k = struct.calcsize(f)
        n = struct.unpack(f, b[i:i + k])[0]
        i += k
        crudo = b[i:i + n]
        return (crudo.decode("utf-8") if t in (0xd9, 0xda, 0xdb) else crudo), i + n
    if t in (0xdc, 0xdd):
        f = ">H" if t == 0xdc else ">I"
        k = struct.calcsize(f)
        return _lista(b, i + k, struct.unpack(f, b[i:i + k])[0])
    if t in (0xde, 0xdf):
        f = ">H" if t == 0xde else ">I"
        k = struct.calcsize(f)
        return _mapa(b, i + k, struct.unpack(f, b[i:i + k])[0])
    if t in (0xd4, 0xd5, 0xd6, 0xd7, 0xd8, 0xc7, 0xc8, 0xc9):     # ext: la fecha (-1) se convierte
        if t in (0xc7, 0xc8, 0xc9):
            f = {0xc7: ">B", 0xc8: ">H", 0xc9: ">I"}[t]
            k = struct.calcsize(f)
            n = struct.unpack(f, b[i:i + k])[0]
            i += k
        else:
            n = {0xd4: 1, 0xd5: 2, 0xd6: 4, 0xd7: 8, 0xd8: 16}[t]
        tipo = struct.unpack("b", b[i:i + 1])[0]
        datos = b[i + 1:i + 1 + n]
        i += 1 + n
        if tipo == -1:
            if n == 4:
                seg, ns = struct.unpack(">I", datos)[0], 0
            elif n == 8:
                v = struct.unpack(">Q", datos)[0]
                seg, ns = v & 0x3ffffffff, v >> 34
            else:
                ns, seg = struct.unpack(">Iq", datos)
            return dt.datetime.fromtimestamp(seg + ns / 1e9, dt.timezone.utc), i
        return datos, i
    raise ValueError(f"tipo MessagePack {t:#x} no soportado")


def _lista(b, i, n):
    salida = []
    for _ in range(n):
        x, i = desempaquetar(b, i)
        salida.append(x)
    return salida, i


def _mapa(b, i, n):
    salida = {}
    for _ in range(n):
        k, i = desempaquetar(b, i)
        v, i = desempaquetar(b, i)
        salida[k] = v
    return salida, i


def marco(msg: bytes) -> bytes:
    """Longitud en varint delante, como pide el protocolo MessagePack de SignalR."""
    n, cab = len(msg), bytearray()
    while True:
        byte = n & 0x7f
        n >>= 7
        cab.append(byte | (0x80 if n else 0))
        if not n:
            return bytes(cab) + msg


def marcos(b: bytes):
    i = 0
    while i < len(b):
        n = desplaza = 0
        while True:
            byte = b[i]
            i += 1
            n |= (byte & 0x7f) << desplaza
            desplaza += 7
            if not byte & 0x80:
                break
        yield desempaquetar(b[i:i + n])[0]
        i += n