Merge branch 'master' of git://github.com/melinath/philo
[philo.git] / models / base.py
index fdfe75c..03b9b54 100644 (file)
@@ -282,34 +282,105 @@ class TreeManager(models.Manager):
        def roots(self):
                return self.filter(parent__isnull=True)
        
-       def get_with_path(self, path, root=None, absolute_result=True, pathsep='/'):
+       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, or None if there is no object with that path,
-               unless absolute_result is set to False, in which case it returns a tuple containing
-               the deepest object found along the path, and the remainder of the path after that
-               object as a string (or None in the case that there is no remaining path).
+               Returns the object with the path, unless absolute_result is set to False, in which
+               case it returns a tuple containing the deepest object found along the path, and the
+               remainder of the path after that object as a string (or None if there is no remaining
+               path). Raises a DoesNotExist exception if no object is found with the given path.
                """
-               slugs = path.split(pathsep)
-               obj = root
-               remaining_slugs = list(slugs)
-               remainder = None
-               for slug in slugs:
-                       remaining_slugs.remove(slug)
-                       if slug: # ignore blank slugs, handles for multiple consecutive pathseps
+               segments = path.split(pathsep)
+               
+               # Check for a trailing pathsep so we can restore it later.
+               trailing_pathsep = False
+               if segments[-1] == '':
+                       trailing_pathsep = True
+               
+               # Clean out blank segments. Handles multiple consecutive pathseps.
+               while True:
+                       try:
+                               segments.remove('')
+                       except ValueError:
+                               break
+               
+               # Special-case a lack of segments. No queries necessary.
+               if not segments:
+                       if root is not None:
+                               return root, None
+                       else:
+                               raise self.model.DoesNotExist('%s matching query does not exist.' % self.model._meta.object_name)
+               
+               def make_query_kwargs(segments):
+                       kwargs = {}
+                       prefix = ""
+                       revsegs = list(segments)
+                       revsegs.reverse()
+                       
+                       for segment in revsegs:
+                               kwargs["%s%s__exact" % (prefix, field)] = segment
+                               prefix += "parent__"
+                       
+                       kwargs[prefix[:-2]] = root
+                       return kwargs
+               
+               def build_path(segments):
+                       path = pathsep.join(segments)
+                       if trailing_pathsep and segments and segments[-1] != '':
+                               path += pathsep
+                       return path
+               
+               def find_obj(segments, depth, deepest_found):
+                       try:
+                               obj = self.get(**make_query_kwargs(segments[:depth]))
+                       except self.model.DoesNotExist:
+                               if absolute_result:
+                                       raise
+                               
+                               depth = (deepest_found + depth)/2
+                               if deepest_found == depth:
+                                       # This should happen if nothing is found with any part of the given path.
+                                       raise
+                               
+                               # Try finding one with half the path since the deepest find.
+                               return find_obj(segments, depth, deepest_found)
+                       else:
+                               # Yay! Found one! Could there be a deeper one?
+                               if absolute_result:
+                                       return obj
+                               
+                               deepest_found = depth
+                               depth = (len(segments) + depth)/2
+                               
+                               if deepest_found == depth:
+                                       return obj, build_path(segments[deepest_found:]) or None
+                               
                                try:
-                                       obj = self.get(slug__exact=slug, parent__exact=obj)
+                                       return find_obj(segments, depth, deepest_found)
                                except self.model.DoesNotExist:
-                                       if absolute_result:
-                                               obj = None
-                                       remaining_slugs.insert(0, slug)
-                                       remainder = pathsep.join(remaining_slugs)
-                                       break
-               if obj:
-                       if absolute_result:
-                               return obj
-                       else:
-                               return (obj, remainder)
-               raise self.model.DoesNotExist('%s matching query does not exist.' % self.model._meta.object_name)
+                                       # Then the deepest one was already found.
+                                       return obj, build_path(segments[deepest_found:])
+               
+               return find_obj(segments, len(segments), 0)
 
 
 class TreeModel(models.Model):
@@ -317,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):