From 5d778648e697651ca681d5347051e6071cfe8487 Mon Sep 17 00:00:00 2001
From: Ross Lawley <ross.lawley@gmail.com>
Date: Fri, 27 May 2011 11:33:40 +0100
Subject: [PATCH] Inital tests for dereferencing improvements

---
 mongoengine/base.py   |   1 +
 mongoengine/fields.py | 215 +++++++++++++++++++++++--------
 mongoengine/tests.py  |  58 +++++++++
 tests/dereference.py  | 288 ++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 512 insertions(+), 50 deletions(-)
 create mode 100644 mongoengine/tests.py
 create mode 100644 tests/dereference.py

diff --git a/mongoengine/base.py b/mongoengine/base.py
index ffceb794..4e3154fd 100644
--- a/mongoengine/base.py
+++ b/mongoengine/base.py
@@ -126,6 +126,7 @@ class BaseField(object):
 
         self.validate(value)
 
+
 class ObjectIdField(BaseField):
     """An field wrapper around MongoDB's ObjectIds.
     """
diff --git a/mongoengine/fields.py b/mongoengine/fields.py
index b2aab5a4..c21829c9 100644
--- a/mongoengine/fields.py
+++ b/mongoengine/fields.py
@@ -337,33 +337,54 @@ class ListField(BaseField):
             # Document class being used rather than a document object
             return self
 
-        if isinstance(self.field, ReferenceField):
-            referenced_type = self.field.document_type
-            # Get value from document instance if available
-            value_list = instance._data.get(self.name)
-            if value_list:
-                deref_list = []
-                for value in value_list:
-                    # Dereference DBRefs
-                    if isinstance(value, (pymongo.dbref.DBRef)):
-                        value = _get_db().dereference(value)
-                        deref_list.append(referenced_type._from_son(value))
-                    else:
-                        deref_list.append(value)
-                instance._data[self.name] = deref_list
+        # Get value from document instance if available
+        value_list = instance._data.get(self.name)
+        if isinstance(self.field, ReferenceField) and value_list:
+            db = _get_db()
+            value_list = [(k,v) for k,v in enumerate(value_list)]
+            deref_list = []
+            collections = {}
 
-        if isinstance(self.field, GenericReferenceField):
-            value_list = instance._data.get(self.name)
-            if value_list:
-                deref_list = []
-                for value in value_list:
-                    # Dereference DBRefs
-                    if isinstance(value, (dict, pymongo.son.SON)):
-                        deref_list.append(self.field.dereference(value))
-                    else:
-                        deref_list.append(value)
-                instance._data[self.name] = deref_list
+            for k, v in value_list:
+                deref_list.append(v)
+                # Save any DBRefs
+                if isinstance(v, (pymongo.dbref.DBRef)):
+                    collections.setdefault(v.collection, []).append((k, v))
 
+            # For each collection get the references
+            for collection, dbrefs in collections.items():
+                id_map = dict([(v.id, k) for k, v in dbrefs])
+                references = db[collection].find({'_id': {'$in': id_map.keys()}})
+                for ref in references:
+                    key = id_map[ref['_id']]
+                    deref_list[key] = get_document(ref['_cls'])._from_son(ref)
+            instance._data[self.name] = deref_list
+
+        # Get value from document instance if available
+        if isinstance(self.field, GenericReferenceField) and value_list:
+
+            db = _get_db()
+            value_list = [(k,v) for k,v in enumerate(value_list)]
+            deref_list = []
+            classes = {}
+
+            for k, v in value_list:
+                deref_list.append(v)
+                # Save any DBRefs
+                if isinstance(v, (dict, pymongo.son.SON)):
+                    classes.setdefault(v['_cls'], []).append((k, v))
+
+            # For each collection get the references
+            for doc_cls, dbrefs in classes.items():
+                id_map = dict([(v['_ref'].id, k) for k, v in dbrefs])
+                doc_cls = get_document(doc_cls)
+                collection = doc_cls._meta['collection']
+                references = db[collection].find({'_id': {'$in': id_map.keys()}})
+
+                for ref in references:
+                    key = id_map[ref['_id']]
+                    deref_list[key] = doc_cls._from_son(ref)
+            instance._data[self.name] = deref_list
         return super(ListField, self).__get__(instance, owner)
 
     def to_python(self, value):
@@ -501,32 +522,53 @@ class MapField(BaseField):
             # Document class being used rather than a document object
             return self
 
-        if isinstance(self.field, ReferenceField):
-            referenced_type = self.field.document_type
-            # Get value from document instance if available
-            value_dict = instance._data.get(self.name)
-            if value_dict:
-                deref_dict = []
-                for key,value in value_dict.iteritems():
-                    # Dereference DBRefs
-                    if isinstance(value, (pymongo.dbref.DBRef)):
-                        value = _get_db().dereference(value)
-                        deref_dict[key] = referenced_type._from_son(value)
-                    else:
-                        deref_dict[key] = value
-                instance._data[self.name] = deref_dict
+        # Get value from document instance if available
+        value_list = instance._data.get(self.name)
+        if isinstance(self.field, ReferenceField) and value_list:
+            db = _get_db()
+            deref_dict = {}
+            collections = {}
 
-        if isinstance(self.field, GenericReferenceField):
-            value_dict = instance._data.get(self.name)
-            if value_dict:
-                deref_dict = []
-                for key,value in value_dict.iteritems():
-                    # Dereference DBRefs
-                    if isinstance(value, (dict, pymongo.son.SON)):
-                        deref_dict[key] = self.field.dereference(value)
-                    else:
-                        deref_dict[key] = value
-                instance._data[self.name] = deref_dict
+            for k, v in value_list.items():
+                deref_dict[k] = v
+                # Save any DBRefs
+                if isinstance(v, (pymongo.dbref.DBRef)):
+                    collections.setdefault(v.collection, []).append((k, v))
+
+            # For each collection get the references
+            for collection, dbrefs in collections.items():
+                id_map = dict([(v.id, k) for k, v in dbrefs])
+                references = db[collection].find({'_id': {'$in': id_map.keys()}})
+                for ref in references:
+                    key = id_map[ref['_id']]
+                    deref_dict[key] = get_document(ref['_cls'])._from_son(ref)
+            instance._data[self.name] = deref_dict
+
+        # Get value from document instance if available
+        if isinstance(self.field, GenericReferenceField) and value_list:
+
+            db = _get_db()
+            value_list = [(k,v) for k,v in value_list.items()]
+            deref_dict = {}
+            classes = {}
+
+            for k, v in value_list:
+                deref_dict[k] = v
+                # Save any DBRefs
+                if isinstance(v, (dict, pymongo.son.SON)):
+                    classes.setdefault(v['_cls'], []).append((k, v))
+
+            # For each collection get the references
+            for doc_cls, dbrefs in classes.items():
+                id_map = dict([(v['_ref'].id, k) for k, v in dbrefs])
+                doc_cls = get_document(doc_cls)
+                collection = doc_cls._meta['collection']
+                references = db[collection].find({'_id': {'$in': id_map.keys()}})
+
+                for ref in references:
+                    key = id_map[ref['_id']]
+                    deref_dict[key] = doc_cls._from_son(ref)
+            instance._data[self.name] = deref_dict
 
         return super(MapField, self).__get__(instance, owner)
 
@@ -869,3 +911,76 @@ class GeoPointField(BaseField):
         if (not isinstance(value[0], (float, int)) and
             not isinstance(value[1], (float, int))):
             raise ValidationError('Both values in point must be float or int.')
+
+
+
+class DereferenceMixin(object):
+    """ WORK IN PROGRESS"""
+
+    def __get__(self, instance, owner):
+        """Descriptor to automatically dereference references.
+        """
+        if instance is None:
+            # Document class being used rather than a document object
+            return self
+
+        # Get value from document instance if available
+        value_list = instance._data.get(self.name)
+        if not value_list:
+            return super(MapField, self).__get__(instance, owner)
+
+        is_dict = True
+        if not hasattr(value_list, 'items'):
+            is_dict = False
+            value_list = dict([(k,v) for k,v in enumerate(value_list)])
+
+        if isinstance(self.field, ReferenceField) and value_list:
+            db = _get_db()
+            dbref = {}
+            if not is_dict:
+                dbref = []
+            collections = {}
+
+            for k, v in value_list.items():
+                dbref[k] = v
+                # Save any DBRefs
+                if isinstance(v, (pymongo.dbref.DBRef)):
+                    collections.setdefault(v.collection, []).append((k, v))
+
+            # For each collection get the references
+            for collection, dbrefs in collections.items():
+                id_map = dict([(v.id, k) for k, v in dbrefs])
+                references = db[collection].find({'_id': {'$in': id_map.keys()}})
+                for ref in references:
+                    key = id_map[ref['_id']]
+                    dbref[key] = get_document(ref['_cls'])._from_son(ref)
+
+            instance._data[self.name] = dbref
+
+        # Get value from document instance if available
+        if isinstance(self.field, GenericReferenceField) and value_list:
+
+            db = _get_db()
+            value_list = [(k,v) for k,v in value_list.items()]
+            dbref = {}
+            classes = {}
+
+            for k, v in value_list:
+                dbref[k] = v
+                # Save any DBRefs
+                if isinstance(v, (dict, pymongo.son.SON)):
+                    classes.setdefault(v['_cls'], []).append((k, v))
+
+            # For each collection get the references
+            for doc_cls, dbrefs in classes.items():
+                id_map = dict([(v['_ref'].id, k) for k, v in dbrefs])
+                doc_cls = get_document(doc_cls)
+                collection = doc_cls._meta['collection']
+                references = db[collection].find({'_id': {'$in': id_map.keys()}})
+
+                for ref in references:
+                    key = id_map[ref['_id']]
+                    dbref[key] = doc_cls._from_son(ref)
+            instance._data[self.name] = dbref
+
+        return super(DereferenceField, self).__get__(instance, owner)
\ No newline at end of file
diff --git a/mongoengine/tests.py b/mongoengine/tests.py
new file mode 100644
index 00000000..4932bc2c
--- /dev/null
+++ b/mongoengine/tests.py
@@ -0,0 +1,58 @@
+from mongoengine.connection import _get_db
+
+class query_counter(object):
+    """ Query_counter contextmanager to get the number of queries. """
+
+    def __init__(self):
+        """ Construct the query_counter. """
+        self.counter = 0
+        self.db = _get_db()
+
+    def __enter__(self):
+        """ On every with block we need to drop the profile collection. """
+        self.db.set_profiling_level(0)
+        self.db.system.profile.drop()
+        self.db.set_profiling_level(2)
+        return self
+
+    def __exit__(self, t, value, traceback):
+        """ Reset the profiling level. """
+        self.db.set_profiling_level(0)
+
+    def __eq__(self, value):
+        """ == Compare querycounter. """
+        return value == self._get_count()
+
+    def __ne__(self, value):
+        """ != Compare querycounter. """
+        return not self.__eq__(value)
+
+    def __lt__(self, value):
+        """ < Compare querycounter. """
+        return self._get_count() < value
+
+    def __le__(self, value):
+        """ <= Compare querycounter. """
+        return self._get_count() <= value
+
+    def __gt__(self, value):
+        """ > Compare querycounter. """
+        return self._get_count() > value
+
+    def __ge__(self, value):
+        """ >= Compare querycounter. """
+        return self._get_count() >= value
+
+    def __int__(self):
+        """ int representation. """
+        return self._get_count()
+
+    def __repr__(self):
+        """ repr query_counter as the number of queries. """
+        return u"%s" % self._get_count()
+
+    def _get_count(self):
+        """ Get the number of queries. """
+        count = self.db.system.profile.find().count() - self.counter
+        self.counter += 1
+        return count
diff --git a/tests/dereference.py b/tests/dereference.py
new file mode 100644
index 00000000..2764ee72
--- /dev/null
+++ b/tests/dereference.py
@@ -0,0 +1,288 @@
+import unittest
+
+from mongoengine import *
+from mongoengine.connection import _get_db
+from mongoengine.tests import query_counter
+
+
+class FieldTest(unittest.TestCase):
+
+    def setUp(self):
+        connect(db='mongoenginetest')
+        self.db = _get_db()
+
+    def ztest_list_item_dereference(self):
+        """Ensure that DBRef items in ListFields are dereferenced.
+        """
+        class User(Document):
+            name = StringField()
+
+        class Group(Document):
+            members = ListField(ReferenceField(User))
+
+        User.drop_collection()
+        Group.drop_collection()
+
+        for i in xrange(1, 51):
+            user = User(name='user %s' % i)
+            user.save()
+
+        group = Group(members=User.objects)
+        group.save()
+
+        with query_counter() as q:
+            self.assertEqual(q, 0)
+
+            group_obj = Group.objects.first()
+            self.assertEqual(q, 1)
+
+            [m for m in group_obj.members]
+            self.assertEqual(q, 2)
+
+        User.drop_collection()
+        Group.drop_collection()
+
+    def ztest_recursive_reference(self):
+        """Ensure that ReferenceFields can reference their own documents.
+        """
+        class Employee(Document):
+            name = StringField()
+            boss = ReferenceField('self')
+            friends = ListField(ReferenceField('self'))
+
+        bill = Employee(name='Bill Lumbergh')
+        bill.save()
+
+        michael = Employee(name='Michael Bolton')
+        michael.save()
+
+        samir = Employee(name='Samir Nagheenanajar')
+        samir.save()
+
+        friends = [michael, samir]
+        peter = Employee(name='Peter Gibbons', boss=bill, friends=friends)
+        peter.save()
+
+        with query_counter() as q:
+            self.assertEqual(q, 0)
+
+            peter = Employee.objects.with_id(peter.id)
+            self.assertEqual(q, 1)
+
+            peter.boss
+            self.assertEqual(q, 2)
+
+            peter.friends
+            self.assertEqual(q, 3)
+
+    def ztest_generic_reference(self):
+
+        class UserA(Document):
+            name = StringField()
+
+        class UserB(Document):
+            name = StringField()
+
+        class UserC(Document):
+            name = StringField()
+
+        class Group(Document):
+            members = ListField(GenericReferenceField())
+
+        UserA.drop_collection()
+        UserB.drop_collection()
+        UserC.drop_collection()
+        Group.drop_collection()
+
+        members = []
+        for i in xrange(1, 51):
+            a = UserA(name='User A %s' % i)
+            a.save()
+
+            b = UserB(name='User B %s' % i)
+            b.save()
+
+            c = UserC(name='User C %s' % i)
+            c.save()
+
+            members += [a, b, c]
+
+        group = Group(members=members)
+        group.save()
+
+        with query_counter() as q:
+            self.assertEqual(q, 0)
+
+            group_obj = Group.objects.first()
+            self.assertEqual(q, 1)
+
+            [m for m in group_obj.members]
+            self.assertEqual(q, 4)
+
+            [m for m in group_obj.members]
+            self.assertEqual(q, 4)
+
+        UserA.drop_collection()
+        UserB.drop_collection()
+        UserC.drop_collection()
+        Group.drop_collection()
+
+    def test_map_field_reference(self):
+
+        class User(Document):
+            name = StringField()
+
+        class Group(Document):
+            members = MapField(ReferenceField(User))
+
+        User.drop_collection()
+        Group.drop_collection()
+
+        members = []
+        for i in xrange(1, 51):
+            user = User(name='user %s' % i)
+            user.save()
+            members.append(user)
+
+        group = Group(members=dict([(str(u.id), u) for u in members]))
+        group.save()
+
+        with query_counter() as q:
+            self.assertEqual(q, 0)
+
+            group_obj = Group.objects.first()
+            self.assertEqual(q, 1)
+
+            [m for m in group_obj.members]
+            self.assertEqual(q, 2)
+
+        User.drop_collection()
+        Group.drop_collection()
+
+    def ztest_generic_reference_dict_field(self):
+
+        class UserA(Document):
+            name = StringField()
+
+        class UserB(Document):
+            name = StringField()
+
+        class UserC(Document):
+            name = StringField()
+
+        class Group(Document):
+            members = DictField()
+
+        UserA.drop_collection()
+        UserB.drop_collection()
+        UserC.drop_collection()
+        Group.drop_collection()
+
+        members = []
+        for i in xrange(1, 51):
+            a = UserA(name='User A %s' % i)
+            a.save()
+
+            b = UserB(name='User B %s' % i)
+            b.save()
+
+            c = UserC(name='User C %s' % i)
+            c.save()
+
+            members += [a, b, c]
+
+        group = Group(members=dict([(str(u.id), u) for u in members]))
+        group.save()
+
+        with query_counter() as q:
+            self.assertEqual(q, 0)
+
+            group_obj = Group.objects.first()
+            self.assertEqual(q, 1)
+
+            [m for m in group_obj.members]
+            self.assertEqual(q, 4)
+
+            [m for m in group_obj.members]
+            self.assertEqual(q, 4)
+
+        group.members = {}
+        group.save()
+
+        with query_counter() as q:
+            self.assertEqual(q, 0)
+
+            group_obj = Group.objects.first()
+            self.assertEqual(q, 1)
+
+            [m for m in group_obj.members]
+            self.assertEqual(q, 1)
+
+        UserA.drop_collection()
+        UserB.drop_collection()
+        UserC.drop_collection()
+        Group.drop_collection()
+
+    def test_generic_reference_map_field(self):
+
+        class UserA(Document):
+            name = StringField()
+
+        class UserB(Document):
+            name = StringField()
+
+        class UserC(Document):
+            name = StringField()
+
+        class Group(Document):
+            members = MapField(GenericReferenceField())
+
+        UserA.drop_collection()
+        UserB.drop_collection()
+        UserC.drop_collection()
+        Group.drop_collection()
+
+        members = []
+        for i in xrange(1, 51):
+            a = UserA(name='User A %s' % i)
+            a.save()
+
+            b = UserB(name='User B %s' % i)
+            b.save()
+
+            c = UserC(name='User C %s' % i)
+            c.save()
+
+            members += [a, b, c]
+
+        group = Group(members=dict([(str(u.id), u) for u in members]))
+        group.save()
+
+        with query_counter() as q:
+            self.assertEqual(q, 0)
+
+            group_obj = Group.objects.first()
+            self.assertEqual(q, 1)
+
+            [m for m in group_obj.members]
+            self.assertEqual(q, 4)
+
+            [m for m in group_obj.members]
+            self.assertEqual(q, 4)
+
+        group.members = {}
+        group.save()
+
+        with query_counter() as q:
+            self.assertEqual(q, 0)
+
+            group_obj = Group.objects.first()
+            self.assertEqual(q, 1)
+
+            [m for m in group_obj.members]
+            self.assertEqual(q, 1)
+
+        UserA.drop_collection()
+        UserB.drop_collection()
+        UserC.drop_collection()
+        Group.drop_collection()
\ No newline at end of file