"""PackStream v1 serialisation, the binary value format used by Bolt.""" import struct from dataclasses import dataclass, field @dataclass class Structure: tag: int fields: list = field(default_factory=list) @dataclass class Node: id: int labels: list properties: dict @dataclass class Relationship: id: int start: int end: int type: str properties: dict @dataclass class Path: nodes: list # alternating sequence is [n0, r0, n1, r1, ..., nk] relationships: list class PackStreamError(Exception): pass class Packer: """Encodes Python values. ``element_ids`` selects the Bolt 5 graph struct layout.""" def __init__(self, element_ids: bool = True): self.element_ids = element_ids self.out = bytearray() def pack(self, value) -> bytes: self.out = bytearray() self._pack(value) return bytes(self.out) def _pack(self, v): out = self.out if v is None: out.append(0xC0) elif v is True: out.append(0xC3) elif v is False: out.append(0xC2) elif isinstance(v, int): self._pack_int(v) elif isinstance(v, float): out.append(0xC1) out += struct.pack(">d", v) elif isinstance(v, str): data = v.encode("utf-8") self._header(len(data), 0x80, 0xD0, 0xD1, 0xD2) out += data elif isinstance(v, (bytes, bytearray)): n = len(v) if n < 0x100: out += bytes([0xCC, n]) elif n < 0x10000: out.append(0xCD) out += struct.pack(">H", n) else: out.append(0xCE) out += struct.pack(">I", n) out += v elif isinstance(v, (list, tuple)): self._header(len(v), 0x90, 0xD4, 0xD5, 0xD6) for item in v: self._pack(item) elif isinstance(v, dict): self._header(len(v), 0xA0, 0xD8, 0xD9, 0xDA) for k, item in v.items(): self._pack(str(k)) self._pack(item) elif isinstance(v, Structure): if len(v.fields) > 15: raise PackStreamError("structure too large") out.append(0xB0 + len(v.fields)) out.append(v.tag) for item in v.fields: self._pack(item) elif isinstance(v, Node): self._pack(self._node(v)) elif isinstance(v, Relationship): fields = [v.id, v.start, v.end, v.type, v.properties] if self.element_ids: fields += [str(v.id), str(v.start), str(v.end)] self._pack(Structure(0x52, fields)) elif isinstance(v, Path): self._pack(self._path(v)) else: raise PackStreamError(f"cannot pack value of type {type(v).__name__}") def _node(self, n: Node) -> Structure: fields = [n.id, n.labels, n.properties] if self.element_ids: fields.append(str(n.id)) return Structure(0x4E, fields) def _path(self, p: Path) -> Structure: nodes, node_index = [], {} for n in p.nodes: if n.id not in node_index: node_index[n.id] = len(nodes) nodes.append(n) rels, rel_index = [], {} for r in p.relationships: if r.id not in rel_index: rel_index[r.id] = len(rels) fields = [r.id, r.type, r.properties] if self.element_ids: fields.append(str(r.id)) rels.append(Structure(0x72, fields)) indices = [] prev = p.nodes[0] for r, n in zip(p.relationships, p.nodes[1:]): idx = rel_index[r.id] + 1 indices.append(idx if r.start == prev.id else -idx) indices.append(node_index[n.id]) prev = n return Structure(0x50, [[self._node(n) for n in nodes], rels, indices]) def _header(self, n, tiny, s8, s16, s32): out = self.out if n < 0x10: out.append(tiny + n) elif n < 0x100: out += bytes([s8, n]) elif n < 0x10000: out.append(s16) out += struct.pack(">H", n) else: out.append(s32) out += struct.pack(">I", n) def _pack_int(self, v): out = self.out if -0x10 <= v < 0x80: out += struct.pack(">b", v) elif -0x80 <= v < 0x80: out.append(0xC8) out += struct.pack(">b", v) elif -0x8000 <= v < 0x8000: out.append(0xC9) out += struct.pack(">h", v) elif -0x80000000 <= v < 0x80000000: out.append(0xCA) out += struct.pack(">i", v) elif -0x8000000000000000 <= v < 0x8000000000000000: out.append(0xCB) out += struct.pack(">q", v) else: raise PackStreamError("integer out of 64-bit range") class Unpacker: def __init__(self, data: bytes): self.data = data self.pos = 0 def _take(self, n): if self.pos + n > len(self.data): raise PackStreamError("unexpected end of data") chunk = self.data[self.pos:self.pos + n] self.pos += n return chunk def _u(self, fmt, n): return struct.unpack(fmt, self._take(n))[0] def unpack(self): marker = self._take(1)[0] if marker < 0x80: return marker if marker >= 0xF0: return marker - 0x100 hi = marker & 0xF0 if hi == 0x80: return self._take(marker & 0x0F).decode("utf-8") if hi == 0x90: return [self.unpack() for _ in range(marker & 0x0F)] if hi == 0xA0: return self._map(marker & 0x0F) if hi == 0xB0: tag = self._take(1)[0] return Structure(tag, [self.unpack() for _ in range(marker & 0x0F)]) if marker == 0xC0: return None if marker == 0xC1: return self._u(">d", 8) if marker == 0xC2: return False if marker == 0xC3: return True if marker == 0xC8: return self._u(">b", 1) if marker == 0xC9: return self._u(">h", 2) if marker == 0xCA: return self._u(">i", 4) if marker == 0xCB: return self._u(">q", 8) if marker in (0xCC, 0xCD, 0xCE): n = {0xCC: lambda: self._u(">B", 1), 0xCD: lambda: self._u(">H", 2), 0xCE: lambda: self._u(">I", 4)}[marker]() return bytes(self._take(n)) if marker in (0xD0, 0xD1, 0xD2): return self._take(self._size(marker - 0xD0)).decode("utf-8") if marker in (0xD4, 0xD5, 0xD6): return [self.unpack() for _ in range(self._size(marker - 0xD4))] if marker in (0xD8, 0xD9, 0xDA): return self._map(self._size(marker - 0xD8)) raise PackStreamError(f"unknown marker 0x{marker:02X}") def _size(self, width_code): return [lambda: self._u(">B", 1), lambda: self._u(">H", 2), lambda: self._u(">I", 4)][width_code]() def _map(self, n): result = {} for _ in range(n): key = self.unpack() result[key] = self.unpack() return result def unpack(data: bytes): return Unpacker(data).unpack()