Commits

Lars Yencken committed 759bf06

Fixes #1, providing a method to automatically number and store an entire tree.

Comments (0)

Files changed (3)

 .DS_Store
 *.egg-info
 __version__.py
+._*
         self.assertEqual(nodes[2], self.fruit)
         pass
 
+    def testWalkMptt(self):
+        nodes = [n for n in self.fruit.walk_mptt()]
+        self.assertEqual(nodes[0], (self.fruit, False))
+        self.assertEqual(nodes[1], (self.fruit.children['lemon'], False))
+        self.assertEqual(nodes[2], (self.fruit.children['lemon'], True))
+        self.assertEqual(nodes[3], (self.fruit.children['strawberry'], False))
+        self.assertEqual(nodes[4], (self.fruit.children['strawberry'], True))
+        self.assertEqual(nodes[5], (self.fruit, True))
+
     def testPrune(self):
         x = self.fruit.copy().prune(lambda n: 'y' not in n.label)
         self.assertEqual(x.label, 'fruit')
 
     def walk(self):
         return self.walk_preorder()
+    
+    def walk_mptt(self):
+        """
+        Walks the tree in a modified pre-order traversal, visiting each node
+        on the way down and again on the way up.
+        """
+        stack = [(self, False)]
+        
+        while stack:
+            node, visited = stack.pop()
+            yield node, visited
+            
+            if not visited:
+                stack.append((node, True))
+                stack.extend((c, False) for c in node.children.values())
 
     def walk_postorder(self):
         """Walks the tree in a post-order manner."""
         self._layout_children(self, '', method)
         return
 
+    def store(self, ModelClass, label_field='label'):
+        """
+        Store this tree to the database using the given Model class. This
+        assumes that every node attribute is usable by the model. Calling this
+        method will also remove the contents of the model.
+        """
+        ModelClass.objects.all().delete()
+        for i, (node, visited) in enumerate(self.walk_mptt()):
+            if visited:
+                node['right_visit'] = i + 1 # 1-indexed
+            else:
+                node['left_visit'] = i + 1 # 1-indexed
+        
+        for node in self.walk():
+            record = node.attrib.copy()
+            record[label_field] = node.label
+            obj = ModelClass(**record)
+            obj.save()
 
     #------------------------------------------------------------------------#
     # PRIVATE