Commits

Gustavo Picon committed 0c3878a

added sorted trees to nestes sets, only move() is missing now

Comments (0)

Files changed (3)

treebeard/models.py

 
 """
 
+import operator
+
+from django.db.models import Q
 from django.db import models, transaction
 
 
         return pos
 
 
+    def get_sorted_pos_queryset(self, siblings, newobj):
+        """
+        :returns: The position a new node will be inserted related to the
+        current node, and also a queryset of the nodes that must be moved
+        to the right. Called only for Node models with :attr:`node_order_by`
+
+        This function was taken from django-mptt (BSD licensed) by Jonathan Buchanan:
+        http://code.google.com/p/django-mptt/source/browse/trunk/mptt/signals.py?spec=svn100&r=100#12
+        """
+
+        fields, filters = [], []
+        for field in self.node_order_by:
+            value = getattr(newobj, field)
+            filters.append(Q(*
+                [Q(**{f: v}) for f, v in fields] +
+                [Q(**{'%s__gt' % field: value})]))
+            fields.append((field, value))
+        return siblings.filter(reduce(operator.or_, filters))
+        try:
+            newpos = self._get_lastpos_in_path(siblings.all()[0].path)
+        except IndexError:
+            newpos, siblings = None, []
+        return newpos, siblings
+
+
     class Meta:
         """
         Abstract model.

treebeard/mp_tree.py

         return newobj
 
 
-    def _get_sorted_pos_queryset(self, siblings, newobj):
-        """
-        :returns: The position a new node will be inserted related to the
-        current node, and also a queryset of the nodes that must be moved
-        to the right. Called only for Node models with :attr:`node_order_by`
-
-        This function was taken from django-mptt (BSD licensed) by Jonathan Buchanan:
-        http://code.google.com/p/django-mptt/source/browse/trunk/mptt/signals.py?spec=svn100&r=100#12
-        """
-
-        fields, filters = [], []
-        for field in self.node_order_by:
-            value = getattr(newobj, field)
-            filters.append(Q(*
-                [Q(**{f: v}) for f, v in fields] +
-                [Q(**{'%s__gt' % field: value})]))
-            fields.append((field, value))
-        siblings = siblings.filter(reduce(operator.or_, filters))
-        try:
-            newpos = self._get_lastpos_in_path(siblings.all()[0].path)
-        except IndexError:
-            newpos, siblings = None, []
-        return newpos, siblings
-
-
     def add_sibling(self, pos=None, **kwargs):
         """
         Adds a new node as a sibling to the current node object.
         newobj.depth = self.depth
         
         if pos == 'sorted-sibling':
-            newpos, siblings = self._get_sorted_pos_queryset(
+            siblings = self.get_sorted_pos_queryset(
                 self.get_siblings(), newobj)
+            try:
+                newpos = self._get_lastpos_in_path(siblings.all()[0].path)
+            except IndexError:
+                newpos = None
             if newpos is None:
                 pos = 'last-sibling'
         else:
             return
         
         if pos == 'sorted-sibling':
-            newpos, siblings = self._get_sorted_pos_queryset(
+            siblings = self.get_sorted_pos_queryset(
                 target.get_siblings(), self)
+            try:
+                newpos = self._get_lastpos_in_path(siblings.all()[0].path)
+            except IndexError:
+                newpos = None
             if newpos is None:
                 pos = 'last-sibling'
 

treebeard/ns_tree.py

             super(NS_NodeQuerySet, self).delete()
             cursor = connection.cursor()
 
-            # Now closing the gap (Celko's trees book, page 62):
+            # Now closing the gap (Celko's trees book, page 62)
             # We do this for every gap that was left in the tree when the nodes
             # were removed.  If many nodes were removed, we're going to update
             # the same nodes over and over again. This would be probably
         """
         return NS_NodeQuerySet(self.model)
 
+
+
 class NS_Node(Node):
     node_order_by = []
 
         newobj = self.__class__(**kwargs)
         newobj.depth = self.depth
 
-        if pos == 'sorted-sibling':
-            raise NotImplemented
         sql = None
+        target = self
 
-        if self.is_root():
+        if target.is_root():
             newobj.lft = 1
             newobj.rgt = 2
-            last_root = self.__class__.get_last_root_node()
+            if pos == 'sorted-sibling':
+                siblings = list(target.get_sorted_pos_queryset(
+                    target.get_siblings(), newobj))
+                if siblings:
+                    pos = 'left'
+                    target = siblings[0]
+                else:
+                    pos = 'last-sibling'
+
+            last_root = target.__class__.get_last_root_node()
             if pos == 'last-sibling' \
-                  or (pos == 'right' and self == last_root):
+                  or (pos == 'right' and target == last_root):
                 newobj.tree_id = last_root.tree_id + 1
             else:
                 newpos = {'first-sibling': 1,
-                          'left': self.tree_id,
-                          'right': self.tree_id + 1}[pos]
+                          'left': target.tree_id,
+                          'right': target.tree_id + 1}[pos]
                 sql = 'UPDATE %(table)s ' \
                       ' SET tree_id = tree_id+1 ' \
                       ' WHERE tree_id >= %(newpos)s' % {
-                          'table': self.__class__._meta.db_table,
+                          'table': target.__class__._meta.db_table,
                           'newpos': newpos
                       }
                 params = []
 
                 newobj.tree_id = newpos
         else:
-            newobj.tree_id = self.tree_id
+            newobj.tree_id = target.tree_id
 
-            target = self
+            if pos == 'sorted-sibling':
+                siblings = list(target.get_sorted_pos_queryset(
+                    target.get_siblings(), newobj))
+                if siblings:
+                    pos = 'left'
+                    target = siblings[0]
+                else:
+                    pos = 'last-sibling'
 
             if pos in ('left', 'right', 'first-sibling'):
-
-                siblings = list(self.get_siblings())
+                siblings = list(target.get_siblings())
 
                 if pos == 'right':
                     if target == siblings[-1]:
                             if found:
                                 target = node
                                 break
-                            elif node == self:
+                            elif node == target:
                                 found = True
                 if pos == 'left':
                     if target == siblings[0]:
                 if pos == 'first-sibling':
                     target = siblings[0]
 
-            move_right = self.__class__._move_right
+            move_right = target.__class__._move_right
 
             if pos == 'last-sibling':
                 newpos = target.get_parent().rgt
-                sql, params = move_right(self.tree_id, newpos, False)
+                sql, params = move_right(target.tree_id, newpos, False)
             elif pos == 'first-sibling':
                 newpos = target.lft
-                sql, params = move_right(self.tree_id, newpos-1, False)
+                sql, params = move_right(target.tree_id, newpos-1, False)
             elif pos == 'left':
                 newpos = target.lft
-                sql, params = move_right(self.tree_id, newpos, True)
+                sql, params = move_right(target.tree_id, newpos, True)
 
             newobj.lft = newpos
             newobj.rgt = newpos + 1
 
         pos = self._fix_move_opts(pos)
 
+        raise NotImplementedError
+
 
     @classmethod
     def load_bulk(cls, bulk_data, parent=None, keep_ids=False):