Quick and dirty initial commit to share progress

d2warehouse/huffman.py

@@ -0,0 +1,57 @@
+from bitarray import bitarray, decodetree
+import itertools
+
+code = {
+    "a": bitarray("11110"),
+    "b": bitarray("0101"),
+    "c": bitarray("01000"),
+    "d": bitarray("110001"),
+    "e": bitarray("110000"),
+    "f": bitarray("010011"),
+    "g": bitarray("11010"),
+    "h": bitarray("00011"),
+    "i": bitarray("1111110"),
+    "j": bitarray("000101110"),
+    "k": bitarray("010010"),
+    "l": bitarray("11101"),
+    "m": bitarray("01101"),
+    "n": bitarray("001101"),
+    "o": bitarray("1111111"),
+    "p": bitarray("11001"),
+    "q": bitarray("11011001"),
+    "r": bitarray("11100"),
+    "s": bitarray("0010"),
+    "t": bitarray("01100"),
+    "u": bitarray("00001"),
+    "v": bitarray("1101110"),
+    "w": bitarray("00000"),
+    "x": bitarray("00111"),
+    "y": bitarray("0001010"),
+    "z": bitarray("11011000"),
+    "0": bitarray("11111011"),
+    "1": bitarray("1111100"),
+    "2": bitarray("001100"),
+    "3": bitarray("1101101"),
+    "4": bitarray("11111010"),
+    "5": bitarray("00010110"),
+    "6": bitarray("1101111"),
+    "7": bitarray("01111"),
+    "8": bitarray("000100"),
+    "9": bitarray("01110"),
+    " ": bitarray("10"),
+}
+
+decode_tree = decodetree(code)
+
+
+def decode(bits: bitarray, n) -> tuple[str, int]:
+    s = "".join(itertools.islice(bits.iterdecode(decode_tree), n))
+    l = len(encode(s))
+    return s, l
+
+
+def encode(s: str) -> bitarray:
+    bits = bitarray(endian="little")
+    bits.encode(code, s)
+    bits.encode(code, " ")
+    return bits