X-Git-Url: http://git.ithinksw.org/philo.git/blobdiff_plain/15c62e5321c25c46ab53154609546fdb05cc5fbe..05bdd3ca32b67c218f7d58c2def2661867b4d1fa:/models/base.py?ds=sidebyside

diff --git a/models/base.py b/models/base.py
index 2338d72..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