Merge branch 'master' of git://github.com/melinath/philo

[philo.git] / models / base.py

diff --git a/models/base.py b/models/base.py
index ae15d16..03b9b54 100644 (file)
--- a/models/base.py
+++ b/models/base.py
@@ -282,6 +282,26 @@ class TreeManager(models.Manager):
        def roots(self):
                return self.filter(parent__isnull=True)
        
+       def get_branch_pks(self, root, depth=5, inclusive=True):
+               branch_pks = []
+               parent_pks = [root.pk]
+               
+               if inclusive:
+                       branch_pks.append(root.pk)
+               
+               for i in xrange(depth):
+                       child_pks = list(self.filter(parent__pk__in=parent_pks).exclude(pk__in=branch_pks).values_list('pk', flat=True))
+                       if not child_pks:
+                               break
+                       
+                       branch_pks += child_pks
+                       parent_pks = child_pks
+               
+               return branch_pks
+       
+       def get_branch(self, root, depth=5, inclusive=True):
+               return self.filter(pk__in=self.get_branch_pks(root, depth, inclusive))
+       
        def get_with_path(self, path, root=None, absolute_result=True, pathsep='/', field='slug'):
                """
                Returns the object with the path, unless absolute_result is set to False, in which
@@ -368,24 +388,47 @@ class TreeModel(models.Model):
        parent = models.ForeignKey('self', related_name='children', null=True, blank=True)
        slug = models.SlugField(max_length=255)
        
-       def has_ancestor(self, ancestor):
-               parent = self
+       def has_ancestor(self, ancestor, inclusive=False):
+               if inclusive:
+                       parent = self
+               else:
+                       parent = self.parent
+               
+               parents = []
+               
                while parent:
                        if parent == ancestor:
                                return True
+                       # If we've found this parent before, the path is recursive and ancestor wasn't on it.
+                       if parent in parents:
+                               return False
+                       parents.append(parent)
                        parent = parent.parent
+               # If ancestor is None, catch it here.
+               if parent == ancestor:
+                       return True
                return False
        
        def get_path(self, root=None, pathsep='/', field='slug'):
-               if root is not None and not self.has_ancestor(root):
-                       raise AncestorDoesNotExist(root)
-               
-               path = getattr(self, field, '?')
                parent = self.parent
+               parents = [self]
+               
+               def compile_path(parents):
+                       return pathsep.join([getattr(parent, field, '?') for parent in parents])
+               
                while parent and parent != root:
-                       path = getattr(parent, field, '?') + pathsep + path
+                       if parent in parents:
+                               if root is not None:
+                                       raise AncestorDoesNotExist(root)
+                               parents.append(parent)
+                               return u"\u2026%s%s" % (pathsep, compile_path(parents[::-1]))
+                       parents.append(parent)
                        parent = parent.parent
-               return path
+               
+               if root is not None and parent is None:
+                       raise AncestorDoesNotExist(root)
+               
+               return compile_path(parents[::-1])
        path = property(get_path)
        
        def __unicode__(self):