X-Git-Url: http://git.ithinksw.org/philo.git/blobdiff_plain/bce7b7c4d0308a1d9a2e450b83ba1dac99e57d5b..05bdd3ca32b67c218f7d58c2def2661867b4d1fa:/models/base.py

diff --git a/models/base.py b/models/base.py
index dabe491..03b9b54 100644
--- 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
@@ -291,6 +311,11 @@ class TreeManager(models.Manager):
 		"""
 		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:
@@ -318,6 +343,12 @@ class TreeManager(models.Manager):
 			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]))
@@ -341,13 +372,13 @@ class TreeManager(models.Manager):
 				depth = (len(segments) + depth)/2
 				
 				if deepest_found == depth:
-					return obj, pathsep.join(segments[deepest_found:]) or None
+					return obj, build_path(segments[deepest_found:]) or None
 				
 				try:
 					return find_obj(segments, depth, deepest_found)
 				except self.model.DoesNotExist:
 					# Then the deepest one was already found.
-					return obj, pathsep.join(segments[deepest_found:])
+					return obj, build_path(segments[deepest_found:])
 		
 		return find_obj(segments, len(segments), 0)
 
@@ -357,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):