draw_tree.py

#!/usr/bin/env python3

import re
import sys

from huffman import Node
from huffman import count_frequencies 
from huffman import build_tree

import pydot


def draw(tree):
    graph = pydot.Dot('tree', graph_type='digraph')
    root = unwind_to(tree, graph)
    print(graph)
    

def unwind_to(tree, graph):
    if not tree:
        return None

    def to_name(c):
        c = ''.join(sorted(c))
        return f'"{c}"'

    def to_label(c, w):
        c = ''.join(sorted(c))
        return f'[{c}] {w}'

    name = to_name(tree[0])
    weight = tree[1]
    left = tree[2]
    right = tree[3]

    node_label = to_label(tree[0], weight)
    new_node = pydot.Node(
            name, label=node_label, shape='box',
            fontname='Fantasque Sans Mono', ordering='in')

    if left:
        left_node = unwind_to(left, graph)
        left_name = to_name(left[0])
        left_weight = left[1]
        left_label = to_label(left[0], left_weight)
        graph.add_edge(pydot.Edge(name, left_name, label=' 0 '))
    if right:
        right_node = unwind_to(right, graph)
        right_name = to_name(right[0])
        right_weight = right[1]
        right_label = to_label(right[0], right_weight)
        graph.add_edge(pydot.Edge(name, right_name, label=' 1 '))

    graph.add_node(new_node)

    return new_node


if __name__ == '__main__':
    if len(sys.argv) < 2:
        dot_cmd = 'dot -Tpng graph.dot > graph.png'
        print(f'usage: {sys.argv[0]} [file] > graph.dot && {dot_cmd}')
        sys.exit(1)
    with open(sys.argv[1]) as f:
        text = f.read()

        # simplification: only a sub-set of characters
        text = re.sub('[^-.,:;!?#/a-zA-Z0-9]', '', text)

        freqs = count_frequencies(text)
        tree = build_tree(freqs)
        draw(tree)