X-Git-Url: http://git.ithinksw.org/philo.git/blobdiff_plain/2f90757a89424650440d684e08cbf53d98cec8fd..d6479a69e07aa3b6756228159d80693aa619b6f8:/models/base.py

diff --git a/models/base.py b/models/base.py
index ae15d16..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
@@ -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):