visualize tree-like text file

plot-decision-tree.py

#!/usr/bin/env python3

import pydotplus
from uuid import uuid4 as uuid
import sys
import re


class Node:
    def __init__(self, attr):
        self.attr = attr
        self.uuid = uuid()
        pass

    def __str__(self):
        return self.attr

    def __id__(self):
        return self.uuid

# for each node (Node, {}) and leaf (Node, {})
# dict key is edge label

'''
sample = (Node("crust size"), {
    "big": (Node("filling size"), {
        "big": (Node("+"), {}),
        "small": (
            Node("shape"), {
                "circle": (Node("+"), {}),
                "triangle/square": (Node("-"), {}),
            }
        )}),
    "small": (Node("shape"), {
        "circle": (Node("+"), {}),
        "square": (Node("filling size"), {
            "big": (Node("+"), {}),
            "small": (Node("-"), {})
        }),
    })})
'''

class Plotter:
    def __init__(self):
        self.graph = pydotplus.Dot(graph_type='digraph')

    def draw(self, parent, child, label):
        p = pydotplus.Node(id(parent), label=str(parent))
        c = pydotplus.Node(id(child), label=str(child))

        self.graph.add_node(p)
        self.graph.add_node(c)

        edge = pydotplus.Edge(p, c, label=label)
        self.graph.add_edge(edge)

    def visit(self, node, parent=None):
        for k, v in node[1].items():
            self.draw(node[0], v[0], k)
            if isinstance(v[1], dict) and v[1]:
                self.visit(v, node[0])

    def plot(self, tree, filename):
        self.visit(tree)
        self.graph.write_png(filename)


class Builder:
    def __init__(self):
        pass

    def build(self, fin):
        root = None
        stack = []

        for line in fin.readlines():

            level = 0
            while line[0] == "\t":
                level += 1
                line = line[1:]

            m = re.search(r'([\w\s\+\-/]+)(\[([\w\s\+\-/]+)\])?', line)
            node = m.group(1).strip()
            try:
                edge = m.group(3).strip()
            except:
                edge = None

            if level > len(stack):
                print("invalid format", file=sys.stderr)
                sys.exit()

            while len(stack) > level:
                stack.pop()

            if level == 0:
                root = [Node(node), {}]
                stack.append(root)
            else:
                tree = [Node(node), {}]
                stack[-1][1].update({edge: tree})
                stack.append(tree)

        return root


if __name__ == "__main__":
    builder = Builder()
    plotter = Plotter()
    plotter.plot(builder.build(sys.stdin), 'decision-tree.png')

test.dat

crust size
	filling size [big]
		+ [big]
		shape [small]
			+ [circle]
			- [triangle/square]
	shape [small]
		+ [circle]
		filling size [square]
			+ [big]
			- [small]