"""Nested Sets"""importsysimportoperatorifsys.version_info>=(3,0):fromfunctoolsimportreducefromdjango.coreimportserializersfromdjango.dbimportconnection,models,transactionfromdjango.db.modelsimportQfromdjango.utils.translationimportugettext_noopas_fromtreebeard.exceptionsimportInvalidMoveToDescendantfromtreebeard.modelsimportNodeclassNS_NodeQuerySet(models.query.QuerySet):""" Custom queryset for the tree node manager. Needed only for the customized delete method. """defdelete(self,removed_ranges=None):""" Custom delete method, will remove all descendant nodes to ensure a consistent tree (no orphans) :returns: ``None`` """ifremoved_rangesisnotNone:# we already know the children, let's call the default django# delete method and let it handle the removal of the user's# foreign keys...super(NS_NodeQuerySet,self).delete()cursor=connection.cursor()# 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# cheaper precalculating the gapsize per intervals, or just do a# complete reordering of the tree (uses COUNT)...fortree_id,drop_lft,drop_rgtinsorted(removed_ranges,reverse=True):sql,params=self.model._get_close_gap_sql(drop_lft,drop_rgt,tree_id)cursor.execute(sql,params)else:# we'll have to manually run through all the nodes that are going# to be deleted and remove nodes from the list if an ancestor is# already getting removed, since that would be redundantremoved={}fornodeinself.order_by('tree_id','lft'):found=Falseforrid,rnodeinremoved.items():ifnode.is_descendant_of(rnode):found=Truebreakifnotfound:removed[node.pk]=node# ok, got the minimal list of nodes to remove...# we must also remove their descendantstoremove=[]ranges=[]forid,nodeinremoved.items():toremove.append(Q(lft__range=(node.lft,node.rgt))&Q(tree_id=node.tree_id))ranges.append((node.tree_id,node.lft,node.rgt))iftoremove:self.model.objects.filter(reduce(operator.or_,toremove)).delete(removed_ranges=ranges)transaction.commit_unless_managed()classNS_NodeManager(models.Manager):""" Custom manager for nodes. """defget_query_set(self):"""Sets the custom queryset as the default."""returnNS_NodeQuerySet(self.model).order_by('tree_id','lft')classNS_Node(Node):"""Abstract model to create your own Nested Sets Trees."""node_order_by=[]lft=models.PositiveIntegerField(db_index=True)rgt=models.PositiveIntegerField(db_index=True)tree_id=models.PositiveIntegerField(db_index=True)depth=models.PositiveIntegerField(db_index=True)objects=NS_NodeManager()@classmethoddefadd_root(cls,**kwargs):"""Adds a root node to the tree."""# do we have a root node already?last_root=cls.get_last_root_node()iflast_rootandlast_root.node_order_by:# there are root nodes and node_order_by has been set# delegate sorted insertion to add_siblingreturnlast_root.add_sibling('sorted-sibling',**kwargs)iflast_root:# adding the new root node as the last onenewtree_id=last_root.tree_id+1else:# adding the first root nodenewtree_id=1# creating the new objectnewobj=cls(**kwargs)newobj.depth=1newobj.tree_id=newtree_idnewobj.lft=1newobj.rgt=2# saving the instance before returning itnewobj.save()transaction.commit_unless_managed()returnnewobj@classmethoddef_move_right(cls,tree_id,rgt,lftmove=False,incdec=2):iflftmove:lftop='>='else:lftop='>'sql='UPDATE %(table)s '\
' SET lft = CASE WHEN lft %(lftop)s%(parent_rgt)d '\
' THEN lft %(incdec)+d '\
' ELSE lft END, '\
' rgt = CASE WHEN rgt >= %(parent_rgt)d '\
' THEN rgt %(incdec)+d '\
' ELSE rgt END '\
' WHERE rgt >= %(parent_rgt)d AND '\
' tree_id = %(tree_id)s'%{'table':connection.ops.quote_name(cls._meta.db_table),'parent_rgt':rgt,'tree_id':tree_id,'lftop':lftop,'incdec':incdec}returnsql,[]@classmethoddef_move_tree_right(cls,tree_id):sql='UPDATE %(table)s '\
' SET tree_id = tree_id+1 '\
' WHERE tree_id >= %(tree_id)d'%{'table':connection.ops.quote_name(cls._meta.db_table),'tree_id':tree_id}returnsql,[]defadd_child(self,**kwargs):"""Adds a child to the node."""ifnotself.is_leaf():# there are child nodes, delegate insertion to add_siblingifself.node_order_by:pos='sorted-sibling'else:pos='last-sibling'last_child=self.get_last_child()last_child._cached_parent_obj=selfreturnlast_child.add_sibling(pos,**kwargs)# we're adding the first child of this nodesql,params=self.__class__._move_right(self.tree_id,self.rgt,False,2)# creating a new objectnewobj=self.__class__(**kwargs)newobj.tree_id=self.tree_idnewobj.depth=self.depth+1newobj.lft=self.lft+1newobj.rgt=self.lft+2# this is just to update the cacheself.rgt+=2newobj._cached_parent_obj=selfcursor=connection.cursor()cursor.execute(sql,params)# saving the instance before returning itnewobj.save()transaction.commit_unless_managed()returnnewobjdefadd_sibling(self,pos=None,**kwargs):"""Adds a new node as a sibling to the current node object."""pos=self._prepare_pos_var_for_add_sibling(pos)# creating a new objectnewobj=self.__class__(**kwargs)newobj.depth=self.depthsql=Nonetarget=selfiftarget.is_root():newobj.lft=1newobj.rgt=2ifpos=='sorted-sibling':siblings=list(target.get_sorted_pos_queryset(target.get_siblings(),newobj))ifsiblings:pos='left'target=siblings[0]else:pos='last-sibling'last_root=target.__class__.get_last_root_node()if((pos=='last-sibling')or(pos=='right'andtarget==last_root)):newobj.tree_id=last_root.tree_id+1else:newpos={'first-sibling':1,'left':target.tree_id,'right':target.tree_id+1}[pos]sql,params=target.__class__._move_tree_right(newpos)newobj.tree_id=newposelse:newobj.tree_id=target.tree_idifpos=='sorted-sibling':siblings=list(target.get_sorted_pos_queryset(target.get_siblings(),newobj))ifsiblings:pos='left'target=siblings[0]else:pos='last-sibling'ifposin('left','right','first-sibling'):siblings=list(target.get_siblings())ifpos=='right':iftarget==siblings[-1]:pos='last-sibling'else:pos='left'found=Falsefornodeinsiblings:iffound:target=nodebreakelifnode==target:found=Trueifpos=='left':iftarget==siblings[0]:pos='first-sibling'ifpos=='first-sibling':target=siblings[0]move_right=self.__class__._move_rightifpos=='last-sibling':newpos=target.get_parent().rgtsql,params=move_right(target.tree_id,newpos,False,2)elifpos=='first-sibling':newpos=target.lftsql,params=move_right(target.tree_id,newpos-1,False,2)elifpos=='left':newpos=target.lftsql,params=move_right(target.tree_id,newpos,True,2)newobj.lft=newposnewobj.rgt=newpos+1# saving the instance before returning itifsql:cursor=connection.cursor()cursor.execute(sql,params)newobj.save()transaction.commit_unless_managed()returnnewobjdefmove(self,target,pos=None):""" Moves the current node and all it's descendants to a new position relative to another node. """pos=self._prepare_pos_var_for_move(pos)cls=self.__class__parent=Noneifposin('first-child','last-child','sorted-child'):# moving to a childiftarget.is_leaf():parent=targetpos='last-child'else:target=target.get_last_child()pos={'first-child':'first-sibling','last-child':'last-sibling','sorted-child':'sorted-sibling'}[pos]iftarget.is_descendant_of(self):raiseInvalidMoveToDescendant(_("Can't move node to a descendant."))ifself==targetand((pos=='left')or(posin('right','last-sibling')andtarget==target.get_last_sibling())or(pos=='first-sibling'andtarget==target.get_first_sibling())):# special cases, not actually moving the node so no need to UPDATEreturnifpos=='sorted-sibling':siblings=list(target.get_sorted_pos_queryset(target.get_siblings(),self))ifsiblings:pos='left'target=siblings[0]else:pos='last-sibling'ifposin('left','right','first-sibling'):siblings=list(target.get_siblings())ifpos=='right':iftarget==siblings[-1]:pos='last-sibling'else:pos='left'found=Falsefornodeinsiblings:iffound:target=nodebreakelifnode==target:found=Trueifpos=='left':iftarget==siblings[0]:pos='first-sibling'ifpos=='first-sibling':target=siblings[0]# ok let's move thiscursor=connection.cursor()move_right=cls._move_rightgap=self.rgt-self.lft+1sql=Nonetarget_tree=target.tree_id# first make a holeifpos=='last-child':newpos=parent.rgtsql,params=move_right(target.tree_id,newpos,False,gap)eliftarget.is_root():newpos=1ifpos=='last-sibling':target_tree=target.get_siblings().reverse()[0].tree_id+1elifpos=='first-sibling':target_tree=1sql,params=cls._move_tree_right(1)elifpos=='left':sql,params=cls._move_tree_right(target.tree_id)else:ifpos=='last-sibling':newpos=target.get_parent().rgtsql,params=move_right(target.tree_id,newpos,False,gap)elifpos=='first-sibling':newpos=target.lftsql,params=move_right(target.tree_id,newpos-1,False,gap)elifpos=='left':newpos=target.lftsql,params=move_right(target.tree_id,newpos,True,gap)ifsql:cursor.execute(sql,params)# we reload 'self' because lft/rgt may have changedfromobj=cls.objects.get(pk=self.pk)depthdiff=target.depth-fromobj.depthifparent:depthdiff+=1# move the tree to the holesql="UPDATE %(table)s "\
" SET tree_id = %(target_tree)d, "\
" lft = lft + %(jump)d , "\
" rgt = rgt + %(jump)d , "\
" depth = depth + %(depthdiff)d "\
" WHERE tree_id = %(from_tree)d AND "\
" lft BETWEEN %(fromlft)d AND %(fromrgt)d"%{'table':connection.ops.quote_name(cls._meta.db_table),'from_tree':fromobj.tree_id,'target_tree':target_tree,'jump':newpos-fromobj.lft,'depthdiff':depthdiff,'fromlft':fromobj.lft,'fromrgt':fromobj.rgt}cursor.execute(sql,[])# close the gapsql,params=cls._get_close_gap_sql(fromobj.lft,fromobj.rgt,fromobj.tree_id)cursor.execute(sql,params)transaction.commit_unless_managed()@classmethoddef_get_close_gap_sql(cls,drop_lft,drop_rgt,tree_id):sql='UPDATE %(table)s '\
' SET lft = CASE '\
' WHEN lft > %(drop_lft)d '\
' THEN lft - %(gapsize)d '\
' ELSE lft END, '\
' rgt = CASE '\
' WHEN rgt > %(drop_lft)d '\
' THEN rgt - %(gapsize)d '\
' ELSE rgt END '\
' WHERE (lft > %(drop_lft)d '\
' OR rgt > %(drop_lft)d) AND '\
' tree_id=%(tree_id)d'%{'table':connection.ops.quote_name(cls._meta.db_table),'gapsize':drop_rgt-drop_lft+1,'drop_lft':drop_lft,'tree_id':tree_id}returnsql,[]@classmethoddefload_bulk(cls,bulk_data,parent=None,keep_ids=False):"""Loads a list/dictionary structure to the tree."""# tree, iterative preorderadded=[]ifparent:parent_id=parent.pkelse:parent_id=None# stack of nodes to analizestack=[(parent_id,node)fornodeinbulk_data[::-1]]whilestack:parent_id,node_struct=stack.pop()# shallow copy of the data strucure so it doesn't persist...node_data=node_struct['data'].copy()ifkeep_ids:node_data['id']=node_struct['id']ifparent_id:parent=cls.objects.get(pk=parent_id)node_obj=parent.add_child(**node_data)else:node_obj=cls.add_root(**node_data)added.append(node_obj.pk)if'children'innode_struct:# extending the stack with the current node as the parent of# the new nodesstack.extend([(node_obj.pk,node)fornodeinnode_struct['children'][::-1]])transaction.commit_unless_managed()returnaddeddefget_children(self):""":returns: A queryset of all the node's children"""returnself.get_descendants().filter(depth=self.depth+1)defget_depth(self):""":returns: the depth (level) of the node"""returnself.depthdefis_leaf(self):""":returns: True if the node is a leaf node (else, returns False)"""returnself.rgt-self.lft==1defget_root(self):""":returns: the root node for the current node object."""ifself.lft==1:returnselfreturnself.__class__.objects.get(tree_id=self.tree_id,lft=1)defis_root(self):""":returns: True if the node is a root node (else, returns False)"""returnself.lft==1defget_siblings(self):""" :returns: A queryset of all the node's siblings, including the node itself. """ifself.lft==1:returnself.get_root_nodes()returnself.get_parent(True).get_children()@classmethoddefdump_bulk(cls,parent=None,keep_ids=True):"""Dumps a tree branch to a python data structure."""qset=cls._get_serializable_model().get_tree(parent)ret,lnk=[],{}forpyobjinqset:serobj=serializers.serialize('python',[pyobj])[0]# django's serializer stores the attributes in 'fields'fields=serobj['fields']depth=fields['depth']# this will be useless in load_bulkdelfields['lft']delfields['rgt']delfields['depth']delfields['tree_id']if'id'infields:# this happens immediately after a load_bulkdelfields['id']newobj={'data':fields}ifkeep_ids:newobj['id']=serobj['pk']if(notparentanddepth==1)or\
(parentanddepth==parent.depth):ret.append(newobj)else:parentobj=pyobj.get_parent()parentser=lnk[parentobj.pk]if'children'notinparentser:parentser['children']=[]parentser['children'].append(newobj)lnk[pyobj.pk]=newobjreturnret@classmethoddefget_tree(cls,parent=None):""" :returns: A *queryset* of nodes ordered as DFS, including the parent. If no parent is given, all trees are returned. """ifparentisNone:# return the entire treereturncls.objects.all()ifparent.is_leaf():returncls.objects.filter(pk=parent.pk)returncls.objects.filter(tree_id=parent.tree_id,lft__range=(parent.lft,parent.rgt-1))defget_descendants(self):""" :returns: A queryset of all the node's descendants as DFS, doesn't include the node itself """ifself.is_leaf():returnself.__class__.objects.none()returnself.__class__.get_tree(self).exclude(pk=self.pk)defget_descendant_count(self):""":returns: the number of descendants of a node."""return(self.rgt-self.lft-1)/2defget_ancestors(self):""" :returns: A queryset containing the current node object's ancestors, starting by the root node and descending to the parent. """ifself.is_root():returnself.__class__.objects.none()returnself.__class__.objects.filter(tree_id=self.tree_id,lft__lt=self.lft,rgt__gt=self.rgt)defis_descendant_of(self,node):""" :returns: ``True`` if the node if a descendant of another node given as an argument, else, returns ``False`` """return(self.tree_id==node.tree_idandself.lft>node.lftandself.rgt<node.rgt)defget_parent(self,update=False):""" :returns: the parent node of the current node object. Caches the result in the object itself to help in loops. """ifself.is_root():returntry:ifupdate:delself._cached_parent_objelse:returnself._cached_parent_objexceptAttributeError:pass# parent = our most direct ancestorself._cached_parent_obj=self.get_ancestors().reverse()[0]returnself._cached_parent_obj@classmethoddefget_root_nodes(cls):""":returns: A queryset containing the root nodes in the tree."""returncls.objects.filter(lft=1)classMeta:"""Abstract model."""abstract=True