From 2e5a27da22a502a56a0e1bf278b906891e502335 Mon Sep 17 00:00:00 2001
From: Stephen Burrows <stephen.r.burrows@gmail.com>
Date: Mon, 28 Mar 2011 17:13:17 -0400
Subject: [PATCH] Implementation of QuerySetMappers for Entities that improves
 the worst-case number of queries required to access all attribute keys.

---
 models/base.py | 53 ++++++++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 51 insertions(+), 2 deletions(-)

diff --git a/models/base.py b/models/base.py
index af1e880..19715ab 100644
--- a/models/base.py
+++ b/models/base.py
@@ -278,6 +278,48 @@ class EntityBase(models.base.ModelBase):
 		return new
 
 
+class EntityAttributeMapper(object, DictMixin):
+	def __init__(self, entity):
+		self.entity = entity
+	
+	def get_attributes(self):
+		return self.entity.attribute_set.all()
+	
+	def make_cache(self):
+		attributes = self.get_attributes()
+		value_lookups = {}
+		
+		for a in attributes:
+			value_lookups.setdefault(a.value_content_type, []).append(a.value_object_id)
+		
+		values_bulk = {}
+		
+		for ct, pks in value_lookups.items():
+			values_bulk[ct] = ct.model_class().objects.in_bulk(pks)
+		
+		self._cache = dict([(a.key, getattr(values_bulk[a.value_content_type].get(a.value_object_id), 'value', None)) for a in attributes])
+	
+	def __getitem__(self, key):
+		if not hasattr(self, '_cache'):
+			self.make_cache()
+		return self._cache[key]
+	
+	def keys(self):
+		if not hasattr(self, '_cache'):
+			self.make_cache()
+		return self._cache.keys()
+	
+	def items(self):
+		if not hasattr(self, '_cache'):
+			self.make_cache()
+		return self._cache.items()
+	
+	def values(self):
+		if not hasattr(self, '_cache'):
+			self.make_cache()
+		return self._cache.values()
+
+
 class Entity(models.Model):
 	__metaclass__ = EntityBase
 	
@@ -285,7 +327,7 @@ class Entity(models.Model):
 	
 	@property
 	def attributes(self):
-		return QuerySetMapper(self.attribute_set.all())
+		return EntityAttributeMapper(self)
 	
 	class Meta:
 		abstract = True
@@ -437,13 +479,20 @@ class TreeEntityBase(MPTTModelBase, EntityBase):
 		return meta.register(cls)
 
 
+class TreeEntityAttributeMapper(EntityAttributeMapper):
+	def get_attributes(self):
+		ancestors = dict(self.entity.get_ancestors(include_self=True).values_list('pk', 'level'))
+		ct = ContentType.objects.get_for_model(self.entity)
+		return sorted(Attribute.objects.filter(entity_content_type=ct, entity_object_id__in=ancestors.keys()), key=lambda x: ancestors[x.entity_object_id])
+
+
 class TreeEntity(Entity, TreeModel):
 	__metaclass__ = TreeEntityBase
 	
 	@property
 	def attributes(self):
 		if self.parent:
-			return QuerySetMapper(self.attribute_set.all(), passthrough=self.parent.attributes)
+			return TreeEntityAttributeMapper(self)
 		return super(TreeEntity, self).attributes
 	
 	class Meta:
-- 
2.20.1