Simple huffman compressor in python inspired by destroyallsoftware.com

jzip.py

import sys
from bitarray import bitarray
from bitstring import BitArray
from bitstring import ConstBitStream
from collections import namedtuple

BUFFER_SIZE = 8192

Node = namedtuple('Node', ['left', 'right', 'count'])
Leaf = namedtuple('Leaf', ['count', 'value'])


def render_tree(tree):
  def _render_tree(tree, label=0):
    if (isinstance(tree, Node)):
      left_label = _render_tree(tree.left, label)
      right_label = _render_tree(tree.right, left_label)
      label = right_label + 1
      print(f'{label} [label="(n = {tree.count})"]\n{label} -> {left_label} [label="0"]\n{label} -> {right_label} [label="1"]')
    else:
      label = label + 1
      print(f'{label} [label="{chr(tree.value)}\\n(n = {tree.count})"]')
    return label
  print('digraph {\n')
  _render_tree(tree)
  print('}\n')

def pack_table(table):
  b = BitArray()
  b.insert(f'uint:8={len(table)}', 0)
  for char, code in table.items():
    # print(f'packing {char} with {code} of length {len(code)}')
    b.insert(f'uint:8={char}', b.len)
    b.insert(f'uint:32={len(code)}', b.len)
    b.insert('0b' + code, b.len)
  return b

def build_table(node, table={}, path=''):
  if (isinstance(node, Node)):
    return { **build_table(node.left, table, path + '0'), **build_table(node.right, table, path + '1') }
  else:
    table[node.value] = path
    return table

def build_tree(data):
  unique = set(data)
  nodes = [Leaf(data.count(c), c) for c in unique]

  # while (leafs):
  while (len(nodes) >= 2):
    nodes.sort(key=lambda a: a.count)
    node1, node2, *rest = nodes
    node = Node(node1, node2, node1.count + node2.count)
    nodes = [*rest, node]

  return nodes[0]

def compress(data, table):
  output = bitarray()
  for c in data:
    output.extend(table[c])
  return output

def unpack_table(data):
  table = {}
  nb_entry = data.read(8).int
  table_length = 8
  for _ in range(nb_entry):
    char = data.read(8).int
    code_length = data.read(32).int
    code = data.read(f'bin:{code_length}')
    table[char] = code
    table_length += 40 + code_length
  return (table_length, table)

def decompress(data):
  (table_length, table) = unpack_table(data)
  data = data[table_length:]
  maxlength = len(max(table.values(), key=len))
  output = []
  while (data):
    for c, code in table.items():
      if data.len >= len(code) and code == data.peek(f'bin:{len(code)}'):
        output += [c]
        data = data[len(code):]
        break
  return output

if __name__ == '__main__':
  if len(sys.argv) == 2 and sys.argv[1] == '-h':
    print('usage: cat <file> | python jzip.py -c')
    print('       cat <zippedfile> | python jzip.py')
    print('       cat file> | python jzip.py -t')
    print('   -c compress standard input')
    print('   -t show the huffman tree based on standard input')
    print(f'limitation: size input limited to {BUFFER_SIZE}')
  elif len(sys.argv) == 2 and sys.argv[1] == '-c':
    input = sys.stdin.buffer.read(BUFFER_SIZE)
    tree = build_tree(input)
    table = build_table(tree)
    output = pack_table(table)
    output.append(compress(input, table))
    sys.stdout.buffer.write(output.tobytes())
  elif len(sys.argv) == 2 and sys.argv[1] == '-t':
    # cat somefile | python jzip.py | dot -Tpng | imgcat
    input = sys.stdin.buffer.read(BUFFER_SIZE)
    tree = build_tree(input)
    render_tree(tree)
  else:
    input = ConstBitStream(BitArray(bytes=sys.stdin.buffer.read(BUFFER_SIZE)))
    sys.stdout.buffer.write(bytearray(decompress(input)))

LINKS.md

http://www.ietf.org/rfc/rfc1950.txt

http://www.ietf.org/rfc/rfc1951.txt

https://zlib.net/feldspar.html

http://www.infinitepartitions.com/art001.html?_sm_nck=1

https://pograph.wordpress.com/2009/06/04/notes-on-gzip-and-deflate-format/

requirements.txt

bpython==0.17.1
bitstring==3.1.5
bitarray==0.8.1

RFC1950 Zkib

Bits identification:
+--------+
|76543210|
+--------+

  0   1   0   1   2   3  
+---+---+---+---+---+---+=====================+---+---+---+---+
|CMF|FLG|     DICTID    |...compressed data...|    ADLER32    |
+---+---+---+---+---+---+=====================+---+---+---+---+
         ‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾         
         optional if
         FLG.FDICT set
CMF:
+--------+
|CMCMINFO|
+--------+
CM (4 bits - 0 to 3): Compression Method
  8: Deflate
CINFO (4 bits - 4 to 7): Compression Info
  If CM == 8, it contains then LZ77 window size is 2^(8 + CINFO)
  CINFO=7 is a window size of 32K.

FLG:
+--------+
|CHECKDLV|
+--------+
CHECK (5 bits - 0 to 4): Check bits for CMF and FLG
  CHECK has a value such as (CMF << 8 | FLG) is a multiple of 31.
FDICT (1 bit - 5): Dictionary bit
  If FDICT == 1 then a dictionary is present. In that case, following the FLG is
  4 bytes which represent the Adler32 of the dictionary.
FLEVEL (2 bits - 6 to 7): Compression Level
  For DEFLATE:
  0 - compressor used fastest algorithm
  1 - compressor used fast algorithm
  2 - compressor used default algorithm
  3 - compressor used maximum compression, slowest algorithm