406 lines
14 KiB
Python
406 lines
14 KiB
Python
|
"""Adjacency List"""
|
||
|
|
|
||
|
|
from django.core import serializers
|
||
|
|
from django.db import models
|
||
|
|
from django.utils.translation import gettext_noop as _
|
||
|
|
from treebeard.exceptions import InvalidMoveToDescendant, NodeAlreadySaved
|
||
|
|
from treebeard.models import Node
|
||
|
|
|
||
|
|
|
||
|
|
def get_result_class(cls):
|
||
|
|
"""
|
||
|
|
For the given model class, determine what class we should use for the
|
||
|
|
nodes returned by its tree methods (such as get_children).
|
||
|
|
|
||
|
|
Usually this will be trivially the same as the initial model class,
|
||
|
|
but there are special cases when model inheritance is in use:
|
||
|
|
|
||
|
|
* If the model extends another via multi-table inheritance, we need to
|
||
|
|
use whichever ancestor originally implemented the tree behaviour (i.e.
|
||
|
|
the one which defines the 'parent' field). We can't use the
|
||
|
|
subclass, because it's not guaranteed that the other nodes reachable
|
||
|
|
from the current one will be instances of the same subclass.
|
||
|
|
|
||
|
|
* If the model is a proxy model, the returned nodes should also use
|
||
|
|
the proxy class.
|
||
|
|
"""
|
||
|
|
base_class = cls._meta.get_field('parent').model
|
||
|
|
if cls._meta.proxy_for_model == base_class:
|
||
|
|
return cls
|
||
|
|
else:
|
||
|
|
return base_class
|
||
|
|
|
||
|
|
|
||
|
|
class AL_NodeManager(models.Manager):
|
||
|
|
"""Custom manager for nodes in an Adjacency List tree."""
|
||
|
|
def get_queryset(self):
|
||
|
|
"""Sets the custom queryset as the default."""
|
||
|
|
if self.model.node_order_by:
|
||
|
|
order_by = ['parent'] + list(self.model.node_order_by)
|
||
|
|
else:
|
||
|
|
order_by = ['parent', 'sib_order']
|
||
|
|
return super().get_queryset().order_by(*order_by)
|
||
|
|
|
||
|
|
|
||
|
|
class AL_Node(Node):
|
||
|
|
"""Abstract model to create your own Adjacency List Trees."""
|
||
|
|
|
||
|
|
objects = AL_NodeManager()
|
||
|
|
node_order_by = None
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def add_root(cls, **kwargs):
|
||
|
|
"""Adds a root node to the tree."""
|
||
|
|
|
||
|
|
if len(kwargs) == 1 and 'instance' in kwargs:
|
||
|
|
# adding the passed (unsaved) instance to the tree
|
||
|
|
newobj = kwargs['instance']
|
||
|
|
if not newobj._state.adding:
|
||
|
|
raise NodeAlreadySaved("Attempted to add a tree node that is "\
|
||
|
|
"already in the database")
|
||
|
|
else:
|
||
|
|
newobj = cls(**kwargs)
|
||
|
|
|
||
|
|
newobj._cached_depth = 1
|
||
|
|
if not cls.node_order_by:
|
||
|
|
try:
|
||
|
|
max = get_result_class(cls).objects.filter(
|
||
|
|
parent__isnull=True).order_by(
|
||
|
|
'sib_order').reverse()[0].sib_order
|
||
|
|
except IndexError:
|
||
|
|
max = 0
|
||
|
|
newobj.sib_order = max + 1
|
||
|
|
newobj.save()
|
||
|
|
return newobj
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def get_root_nodes(cls):
|
||
|
|
""":returns: A queryset containing the root nodes in the tree."""
|
||
|
|
return get_result_class(cls).objects.filter(parent__isnull=True)
|
||
|
|
|
||
|
|
def get_depth(self, update=False):
|
||
|
|
"""
|
||
|
|
:returns: the depth (level) of the node
|
||
|
|
Caches the result in the object itself to help in loops.
|
||
|
|
|
||
|
|
:param update: Updates the cached value.
|
||
|
|
"""
|
||
|
|
|
||
|
|
if self.parent_id is None:
|
||
|
|
return 1
|
||
|
|
|
||
|
|
try:
|
||
|
|
if update:
|
||
|
|
del self._cached_depth
|
||
|
|
else:
|
||
|
|
return self._cached_depth
|
||
|
|
except AttributeError:
|
||
|
|
pass
|
||
|
|
|
||
|
|
depth = 0
|
||
|
|
node = self
|
||
|
|
while node:
|
||
|
|
node = node.parent
|
||
|
|
depth += 1
|
||
|
|
self._cached_depth = depth
|
||
|
|
return depth
|
||
|
|
|
||
|
|
def get_children(self):
|
||
|
|
""":returns: A queryset of all the node's children"""
|
||
|
|
return get_result_class(self.__class__).objects.filter(parent=self)
|
||
|
|
|
||
|
|
def get_parent(self, update=False):
|
||
|
|
""":returns: the parent node of the current node object."""
|
||
|
|
if self._meta.proxy_for_model:
|
||
|
|
# the current node is a proxy model; the returned parent
|
||
|
|
# should be the same proxy model, so we need to explicitly
|
||
|
|
# fetch it as an instance of that model rather than simply
|
||
|
|
# following the 'parent' relation
|
||
|
|
if self.parent_id is None:
|
||
|
|
return None
|
||
|
|
else:
|
||
|
|
return self.__class__.objects.get(pk=self.parent_id)
|
||
|
|
else:
|
||
|
|
return self.parent
|
||
|
|
|
||
|
|
def get_ancestors(self):
|
||
|
|
"""
|
||
|
|
:returns: A *list* containing the current node object's ancestors,
|
||
|
|
starting by the root node and descending to the parent.
|
||
|
|
"""
|
||
|
|
ancestors = []
|
||
|
|
if self._meta.proxy_for_model:
|
||
|
|
# the current node is a proxy model; our result set
|
||
|
|
# should use the same proxy model, so we need to
|
||
|
|
# explicitly fetch instances of that model
|
||
|
|
# when following the 'parent' relation
|
||
|
|
cls = self.__class__
|
||
|
|
node = self
|
||
|
|
while node.parent_id:
|
||
|
|
node = cls.objects.get(pk=node.parent_id)
|
||
|
|
ancestors.insert(0, node)
|
||
|
|
else:
|
||
|
|
node = self.parent
|
||
|
|
while node:
|
||
|
|
ancestors.insert(0, node)
|
||
|
|
node = node.parent
|
||
|
|
return ancestors
|
||
|
|
|
||
|
|
def get_root(self):
|
||
|
|
""":returns: the root node for the current node object."""
|
||
|
|
ancestors = self.get_ancestors()
|
||
|
|
if ancestors:
|
||
|
|
return ancestors[0]
|
||
|
|
return self
|
||
|
|
|
||
|
|
def is_descendant_of(self, node):
|
||
|
|
"""
|
||
|
|
:returns: ``True`` if the node if a descendant of another node given
|
||
|
|
as an argument, else, returns ``False``
|
||
|
|
"""
|
||
|
|
return self.pk in [obj.pk for obj in node.get_descendants()]
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def dump_bulk(cls, parent=None, keep_ids=True):
|
||
|
|
"""Dumps a tree branch to a python data structure."""
|
||
|
|
|
||
|
|
serializable_cls = cls._get_serializable_model()
|
||
|
|
if (
|
||
|
|
parent and serializable_cls != cls and
|
||
|
|
parent.__class__ != serializable_cls
|
||
|
|
):
|
||
|
|
parent = serializable_cls.objects.get(pk=parent.pk)
|
||
|
|
|
||
|
|
# a list of nodes: not really a queryset, but it works
|
||
|
|
objs = serializable_cls.get_tree(parent)
|
||
|
|
|
||
|
|
ret, lnk = [], {}
|
||
|
|
pk_field = cls._meta.pk.attname
|
||
|
|
for node, pyobj in zip(objs, serializers.serialize('python', objs)):
|
||
|
|
depth = node.get_depth()
|
||
|
|
# django's serializer stores the attributes in 'fields'
|
||
|
|
fields = pyobj['fields']
|
||
|
|
del fields['parent']
|
||
|
|
|
||
|
|
# non-sorted trees have this
|
||
|
|
if 'sib_order' in fields:
|
||
|
|
del fields['sib_order']
|
||
|
|
|
||
|
|
if pk_field in fields:
|
||
|
|
del fields[pk_field]
|
||
|
|
|
||
|
|
newobj = {'data': fields}
|
||
|
|
if keep_ids:
|
||
|
|
newobj[pk_field] = pyobj['pk']
|
||
|
|
|
||
|
|
if (not parent and depth == 1) or\
|
||
|
|
(parent and depth == parent.get_depth()):
|
||
|
|
ret.append(newobj)
|
||
|
|
else:
|
||
|
|
parentobj = lnk[node.parent_id]
|
||
|
|
if 'children' not in parentobj:
|
||
|
|
parentobj['children'] = []
|
||
|
|
parentobj['children'].append(newobj)
|
||
|
|
lnk[node.pk] = newobj
|
||
|
|
return ret
|
||
|
|
|
||
|
|
def add_child(self, **kwargs):
|
||
|
|
"""Adds a child to the node."""
|
||
|
|
cls = get_result_class(self.__class__)
|
||
|
|
|
||
|
|
if len(kwargs) == 1 and 'instance' in kwargs:
|
||
|
|
# adding the passed (unsaved) instance to the tree
|
||
|
|
newobj = kwargs['instance']
|
||
|
|
if not newobj._state.adding:
|
||
|
|
raise NodeAlreadySaved("Attempted to add a tree node that is "\
|
||
|
|
"already in the database")
|
||
|
|
else:
|
||
|
|
newobj = cls(**kwargs)
|
||
|
|
|
||
|
|
try:
|
||
|
|
newobj._cached_depth = self._cached_depth + 1
|
||
|
|
except AttributeError:
|
||
|
|
pass
|
||
|
|
if not cls.node_order_by:
|
||
|
|
try:
|
||
|
|
max = cls.objects.filter(parent=self).reverse(
|
||
|
|
)[0].sib_order
|
||
|
|
except IndexError:
|
||
|
|
max = 0
|
||
|
|
newobj.sib_order = max + 1
|
||
|
|
newobj.parent = self
|
||
|
|
newobj.save()
|
||
|
|
return newobj
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def _get_tree_recursively(cls, results, parent, depth):
|
||
|
|
if parent:
|
||
|
|
nodes = parent.get_children()
|
||
|
|
else:
|
||
|
|
nodes = cls.get_root_nodes()
|
||
|
|
for node in nodes:
|
||
|
|
node._cached_depth = depth
|
||
|
|
results.append(node)
|
||
|
|
cls._get_tree_recursively(results, node, depth + 1)
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def get_tree(cls, parent=None):
|
||
|
|
"""
|
||
|
|
:returns: A list of nodes ordered as DFS, including the parent. If
|
||
|
|
no parent is given, the entire tree is returned.
|
||
|
|
"""
|
||
|
|
if parent:
|
||
|
|
depth = parent.get_depth() + 1
|
||
|
|
results = [parent]
|
||
|
|
else:
|
||
|
|
depth = 1
|
||
|
|
results = []
|
||
|
|
cls._get_tree_recursively(results, parent, depth)
|
||
|
|
return results
|
||
|
|
|
||
|
|
def get_descendants(self):
|
||
|
|
"""
|
||
|
|
:returns: A *list* of all the node's descendants, doesn't
|
||
|
|
include the node itself
|
||
|
|
"""
|
||
|
|
return self.__class__.get_tree(parent=self)[1:]
|
||
|
|
|
||
|
|
def get_descendant_count(self):
|
||
|
|
""":returns: the number of descendants of a nodee"""
|
||
|
|
return len(self.get_descendants())
|
||
|
|
|
||
|
|
def get_siblings(self):
|
||
|
|
"""
|
||
|
|
:returns: A queryset of all the node's siblings, including the node
|
||
|
|
itself.
|
||
|
|
"""
|
||
|
|
if self.parent:
|
||
|
|
return get_result_class(self.__class__).objects.filter(
|
||
|
|
parent=self.parent)
|
||
|
|
return self.__class__.get_root_nodes()
|
||
|
|
|
||
|
|
def add_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)
|
||
|
|
|
||
|
|
if len(kwargs) == 1 and 'instance' in kwargs:
|
||
|
|
# adding the passed (unsaved) instance to the tree
|
||
|
|
newobj = kwargs['instance']
|
||
|
|
if not newobj._state.adding:
|
||
|
|
raise NodeAlreadySaved("Attempted to add a tree node that is "\
|
||
|
|
"already in the database")
|
||
|
|
else:
|
||
|
|
# creating a new object
|
||
|
|
newobj = get_result_class(self.__class__)(**kwargs)
|
||
|
|
|
||
|
|
if not self.node_order_by:
|
||
|
|
newobj.sib_order = self.__class__._get_new_sibling_order(pos,
|
||
|
|
self)
|
||
|
|
newobj.parent_id = self.parent_id
|
||
|
|
newobj.save()
|
||
|
|
return newobj
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def _is_target_pos_the_last_sibling(cls, pos, target):
|
||
|
|
return pos == 'last-sibling' or (
|
||
|
|
pos == 'right' and target == target.get_last_sibling())
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def _make_hole_in_db(cls, min, target_node):
|
||
|
|
qset = get_result_class(cls).objects.filter(sib_order__gte=min)
|
||
|
|
if target_node.is_root():
|
||
|
|
qset = qset.filter(parent__isnull=True)
|
||
|
|
else:
|
||
|
|
qset = qset.filter(parent=target_node.parent)
|
||
|
|
qset.update(sib_order=models.F('sib_order') + 1)
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def _make_hole_and_get_sibling_order(cls, pos, target_node):
|
||
|
|
siblings = target_node.get_siblings()
|
||
|
|
siblings = {
|
||
|
|
'left': siblings.filter(sib_order__gte=target_node.sib_order),
|
||
|
|
'right': siblings.filter(sib_order__gt=target_node.sib_order),
|
||
|
|
'first-sibling': siblings
|
||
|
|
}[pos]
|
||
|
|
sib_order = {
|
||
|
|
'left': target_node.sib_order,
|
||
|
|
'right': target_node.sib_order + 1,
|
||
|
|
'first-sibling': 1
|
||
|
|
}[pos]
|
||
|
|
try:
|
||
|
|
min = siblings.order_by('sib_order')[0].sib_order
|
||
|
|
except IndexError:
|
||
|
|
min = 0
|
||
|
|
if min:
|
||
|
|
cls._make_hole_in_db(min, target_node)
|
||
|
|
return sib_order
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def _get_new_sibling_order(cls, pos, target_node):
|
||
|
|
if cls._is_target_pos_the_last_sibling(pos, target_node):
|
||
|
|
sib_order = target_node.get_last_sibling().sib_order + 1
|
||
|
|
else:
|
||
|
|
sib_order = cls._make_hole_and_get_sibling_order(pos, target_node)
|
||
|
|
return sib_order
|
||
|
|
|
||
|
|
def move(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)
|
||
|
|
|
||
|
|
sib_order = None
|
||
|
|
parent = None
|
||
|
|
|
||
|
|
if pos in ('first-child', 'last-child', 'sorted-child'):
|
||
|
|
# moving to a child
|
||
|
|
if not target.is_leaf():
|
||
|
|
target = target.get_last_child()
|
||
|
|
pos = {'first-child': 'first-sibling',
|
||
|
|
'last-child': 'last-sibling',
|
||
|
|
'sorted-child': 'sorted-sibling'}[pos]
|
||
|
|
else:
|
||
|
|
parent = target
|
||
|
|
if pos == 'sorted-child':
|
||
|
|
pos = 'sorted-sibling'
|
||
|
|
else:
|
||
|
|
pos = 'first-sibling'
|
||
|
|
sib_order = 1
|
||
|
|
|
||
|
|
if target.is_descendant_of(self):
|
||
|
|
raise InvalidMoveToDescendant(
|
||
|
|
_("Can't move node to a descendant."))
|
||
|
|
|
||
|
|
if self == target and (
|
||
|
|
(pos == 'left') or
|
||
|
|
(pos in ('right', 'last-sibling') and
|
||
|
|
target == target.get_last_sibling()) or
|
||
|
|
(pos == 'first-sibling' and
|
||
|
|
target == target.get_first_sibling())):
|
||
|
|
# special cases, not actually moving the node so no need to UPDATE
|
||
|
|
return
|
||
|
|
|
||
|
|
if pos == 'sorted-sibling':
|
||
|
|
if parent:
|
||
|
|
self.parent = parent
|
||
|
|
else:
|
||
|
|
self.parent = target.parent
|
||
|
|
else:
|
||
|
|
if sib_order:
|
||
|
|
self.sib_order = sib_order
|
||
|
|
else:
|
||
|
|
self.sib_order = self.__class__._get_new_sibling_order(pos,
|
||
|
|
target)
|
||
|
|
if parent:
|
||
|
|
self.parent = parent
|
||
|
|
else:
|
||
|
|
self.parent = target.parent
|
||
|
|
|
||
|
|
self.save()
|
||
|
|
|
||
|
|
class Meta:
|
||
|
|
"""Abstract model."""
|
||
|
|
abstract = True
|