Added no_dereference context manager (#82)

Reorganised the context_managers as well

tests/queryset/queryset.py

@@ -17,7 +17,7 @@ from bson import ObjectId
 from mongoengine import *
 from mongoengine.connection import get_connection
 from mongoengine.python_support import PY3
-from mongoengine.tests import query_counter
+from mongoengine.context_managers import query_counter
 from mongoengine.queryset import (QuerySet, QuerySetManager,
                                   MultipleObjectsReturned, DoesNotExist,
                                   QueryFieldList, queryset_manager)

tests/test_connection.py

@@ -1,3 +1,4 @@
+from __future__ import with_statement
 import datetime
 import pymongo
 import unittest
@@ -8,6 +9,7 @@ from bson.tz_util import utc
 
 from mongoengine import *
 from mongoengine.connection import get_db, get_connection, ConnectionError
+from mongoengine.context_managers import switch_db
 
 
 class ConnectionTest(unittest.TestCase):
@@ -105,7 +107,7 @@ class ConnectionTest(unittest.TestCase):
         Group(name="hello - default").save()
         self.assertEqual(1, Group.objects.count())
 
-        with SwitchDB(Group, 'testdb-1') as Group:
+        with switch_db(Group, 'testdb-1') as Group:
 
             self.assertEqual(0, Group.objects.count())
 

tests/test_dereference.py

@@ -8,7 +8,7 @@ from bson import DBRef, ObjectId
 
 from mongoengine import *
 from mongoengine.connection import get_db
-from mongoengine.tests import query_counter
+from mongoengine.context_managers import query_counter, no_dereference
 
 
 class FieldTest(unittest.TestCase):
@@ -1121,5 +1121,77 @@ class FieldTest(unittest.TestCase):
 
             self.assertEqual(q, 2)
 
+    def test_no_dereference_context_manager_object_id(self):
+        """Ensure that DBRef items in ListFields aren't dereferenced.
+        """
+        class User(Document):
+            name = StringField()
+
+        class Group(Document):
+            ref = ReferenceField(User, dbref=False)
+            generic = GenericReferenceField()
+            members = ListField(ReferenceField(User, dbref=False))
+
+        User.drop_collection()
+        Group.drop_collection()
+
+        for i in xrange(1, 51):
+            User(name='user %s' % i).save()
+
+        user = User.objects.first()
+        Group(ref=user, members=User.objects, generic=user).save()
+
+        with no_dereference(Group) as NoDeRefGroup:
+            self.assertTrue(Group._fields['members']._auto_dereference)
+            self.assertFalse(NoDeRefGroup._fields['members']._auto_dereference)
+
+        with no_dereference(Group) as Group:
+            group = Group.objects.first()
+            self.assertTrue(all([not isinstance(m, User)
+                                for m in group.members]))
+            self.assertFalse(isinstance(group.ref, User))
+            self.assertFalse(isinstance(group.generic, User))
+
+        self.assertTrue(all([isinstance(m, User)
+                             for m in group.members]))
+        self.assertTrue(isinstance(group.ref, User))
+        self.assertTrue(isinstance(group.generic, User))
+
+    def test_no_dereference_context_manager_dbref(self):
+        """Ensure that DBRef items in ListFields aren't dereferenced.
+        """
+        class User(Document):
+            name = StringField()
+
+        class Group(Document):
+            ref = ReferenceField(User, dbref=True)
+            generic = GenericReferenceField()
+            members = ListField(ReferenceField(User, dbref=True))
+
+        User.drop_collection()
+        Group.drop_collection()
+
+        for i in xrange(1, 51):
+            User(name='user %s' % i).save()
+
+        user = User.objects.first()
+        Group(ref=user, members=User.objects, generic=user).save()
+
+        with no_dereference(Group) as NoDeRefGroup:
+            self.assertTrue(Group._fields['members']._auto_dereference)
+            self.assertFalse(NoDeRefGroup._fields['members']._auto_dereference)
+
+        with no_dereference(Group) as Group:
+            group = Group.objects.first()
+            self.assertTrue(all([not isinstance(m, User)
+                                for m in group.members]))
+            self.assertFalse(isinstance(group.ref, User))
+            self.assertFalse(isinstance(group.generic, User))
+
+        self.assertTrue(all([isinstance(m, User)
+                             for m in group.members]))
+        self.assertTrue(isinstance(group.ref, User))
+        self.assertTrue(isinstance(group.generic, User))
+
 if __name__ == '__main__':
     unittest.main()