Easy way of MPTT Fixtures: parsing tree and creating objects manually

mptt-load

"""
Parses a given file as a category tree and saves each item as MPTT Objects.
Example format of the input file:

- Books
-- Novel
-- Tech
--- Science-Fiction
--- Documentary
-- Fun
- Electronics
-- PC
-- Tablets
-- Phones
--- Smart Phones
--- Touch Phones
--- Cell Phones
-- TVs
-Other Stuff
"""

from models import Category

MODEL = Category

class Node:
    def __init__(self, name):
        self.name = name
        self.children = []

    def add_child(self, node):
        self.children.append(node)

    def __str__(self):
        return self.name


def set_defaults(obj):
    """Put whatever extra attribute you'd like to assign"""
    obj.site_id = 1

class Tree:
    def __init__(self):
        self.root = Node("Root")

    def print_node(self, node, base_hyphen):
        print "%s%s" % (base_hyphen*"\t",  node.name)
        for child in node.children:
            self.print_node(child, base_hyphen+1)

    def print_tree(self):
        for node in self.root.children:
            self.print_node(node, 0)

    def save_node(self, node, parent=None):
        obj = MODEL()
        obj.name = node.name
        obj.parent = parent
        set_defaults(obj)
        obj.save()

        for child in node.children:
            self.save_node(child, obj)

    def save_tree(self):
        for node in self.root.children:
            self.save_node(node)



class TreeParser:
    def __init__(self, filename):
        self.text = open(filename).read()
        self.tree = Tree()
        self.hyphen_dict = {}

    def parse(self):
        """
        if hyphen (-) count is 1 (top level)
            Add to root's children
        else:
            Find latest item with hyphen-1 hyphenation and assign yourself to that item.
        """

        for line in self.text.split("\n"):
            # Get the hyphen count.
            cols = line.split(" ")
            hyphen_count = cols[0].count("-")
            text = " ".join(cols[1:])
            node = Node(text.strip())

            if hyphen_count == 1:
                # if it's top level, add to root's children.
                self.tree.root.add_child(node)
                self.hyphen_dict[1] = node
            else:
                # if it's not top level, Find latest item with hyphen-1 hyphenation
                # and assign yourself to that item.
                try:
                    last_parent = self.hyphen_dict[hyphen_count-1]
                    last_parent.add_child(node)
                    self.hyphen_dict[hyphen_count] = node
                except KeyError:
                    print "Invalid tree structure. Check hyphens."

        self.tree.print_tree()
        self.tree.save_tree()


if __name__ == "__main__":
    parser = TreeParser("tree_data.txt")
    parser.parse()