Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
Easy way of MPTT Fixtures: parsing tree and creating objects manually
"""
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()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment