From: Stephen Burrows <stephen.r.burrows@gmail.com>
Date: Wed, 10 Nov 2010 19:07:42 +0000 (-0500)
Subject: Added get_branch support to TreeManager.
X-Git-Tag: philo-0.9~26^2~3
X-Git-Url: http://git.ithinksw.org/philo.git/commitdiff_plain/d6479a69e07aa3b6756228159d80693aa619b6f8

Added get_branch support to TreeManager.
---

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