diff --git a/erldistpy/etf.py b/erldistpy/etf.py index 9b502d8..765b5f5 100644 --- a/erldistpy/etf.py +++ b/erldistpy/etf.py @@ -27,6 +27,7 @@ LIST_EXT = 108 BINARY_EXT = 109 SMALL_BIG_EXT = 110 LARGE_BIG_EXT = 111 +MAP_EXT = 116 NEW_PID_EXT = 88 NEWER_REFERENCE_EXT = 90 ATOM_UTF8_EXT = 118 @@ -90,6 +91,8 @@ def _encode(term: object) -> bytes: return _encode_tuple(term) if isinstance(term, list): return _encode_list(term) + if isinstance(term, dict): + return _encode_map(term) if isinstance(term, Pid): return _encode_pid(term) if isinstance(term, Reference): @@ -134,6 +137,12 @@ def _encode_tuple(t: tuple) -> bytes: return out + b"".join(_encode(x) for x in t) +def _encode_map(d: dict) -> bytes: + out = bytes([MAP_EXT]) + struct.pack(">I", len(d)) + out += b"".join(_encode(k) + _encode(v) for k, v in d.items()) + return out + + def _encode_list(lst: list) -> bytes: if not lst: return bytes([NIL_EXT]) @@ -249,6 +258,15 @@ def _decode(data: bytes, off: int) -> tuple[object, int]: # Improper list — uncommon in our protocol; surface it raise ETFError(f"improper list tail: {tail!r}") return items, off + if tag == MAP_EXT: + arity = struct.unpack(">I", data[off:off + 4])[0] + off += 4 + result: dict = {} + for _ in range(arity): + k, off = _decode(data, off) + v, off = _decode(data, off) + result[k] = v + return result, off if tag == NEW_PID_EXT: node, off = _decode(data, off) if not isinstance(node, Atom): diff --git a/tests/test_etf.py b/tests/test_etf.py index 9eccc30..b43b530 100644 --- a/tests/test_etf.py +++ b/tests/test_etf.py @@ -49,6 +49,13 @@ GOLDEN = [ ("nested", "83680364000867656e5f63616c6c6c0000000168026400046e6f646564000c77616c6c65744063616d6d796a680264000b6765745f62616c616e6365640003786d72", (Atom("gen_call"), [(Atom("node"), Atom("wallet@cammy"))], (Atom("get_balance"), Atom("xmr")))), + ("empty_map", "837400000000", {}), + ("addr_map", + "8374000000026d00000007616464726573736d00000005786d722d316d0000000863757272656e6379640003786d72", + {b"address": b"xmr-1", b"currency": Atom("xmr")}), + ("nested_ok", + "8368026400026f6b740000000264000762616c616e6365620012d6876400067374617475736400026f6b", + (Atom("ok"), {Atom("balance"): 1234567, Atom("status"): Atom("ok")})), ] @@ -67,6 +74,8 @@ def test_decode_golden(name, hex_in, expected): (), (1,), (Atom("ok"), b"value"), [], [1, 2, 3], [Atom("a"), b"b", 3], (Atom("gen_call"), [], (Atom("ping"),)), + {}, {Atom("k"): 1}, {b"a": b"b", b"c": 42}, + (Atom("ok"), {Atom("status"): Atom("ok"), Atom("count"): 7}), ]) def test_round_trip(term): assert decode(encode(term)) == term