Removed CollectionManager, moved work to QuerySet

As CollectionManager has been replaced with QuerySet and
QuerySetManager, collection.py has been renamed queryset.py.

mongoengine/base.py

@@ -1,4 +1,4 @@
-from collection import CollectionManager
+from queryset import QuerySetManager
 
 import pymongo
 
@@ -153,7 +153,7 @@ class TopLevelDocumentMetaclass(DocumentMetaclass):
         # Set up collection manager, needs the class to have fields so use
         # DocumentMetaclass before instantiating CollectionManager object
         new_class = super_new(cls, name, bases, attrs)
-        new_class.objects = CollectionManager(new_class)
+        new_class.objects = QuerySetManager(new_class)
 
         return new_class
 

mongoengine/document.py

@@ -18,7 +18,8 @@ class Document(BaseDocument):
         """Save the document to the database. If the document already exists,
         it will be updated, otherwise it will be created.
         """
-        self.objects._save_document(self)
+        _id = self.objects._collection.save(self._to_mongo())
+        self._id = _id
 
     @classmethod
     def drop_collection(cls):

mongoengine/fields.py

@@ -145,7 +145,8 @@ class ReferenceField(BaseField):
         # Dereference DBRefs
         if isinstance(value, (pymongo.dbref.DBRef)):
             value = _get_db().dereference(value)
-            instance._data[self.name] = self.document_type._from_son(value)
+            if value is not None:
+                instance._data[self.name] = self.document_type._from_son(value)
         
         return super(ReferenceField, self).__get__(instance, owner)
 

mongoengine/queryset.py

@@ -8,14 +8,18 @@ class QuerySet(object):
     providing Document objects as the results.
     """
     
-    def __init__(self, document, collection, query):
+    def __init__(self, document, collection):
         self._document = document
         self._collection = collection
-
+        self._query = {'_types': self._document._class_name}
-        self._query = QuerySet._transform_query(**query)
-        self._query['_types'] = self._document._class_name
         self._cursor_obj = None
 
+    def __call__(self, **query):
+        """Filter the selected documents by calling the queryset with a query.
+        """
+        self._query.update(QuerySet._transform_query(**query))
+        return self
+
     @property
     def _cursor(self):
         if not self._cursor_obj:
@@ -42,6 +46,25 @@ class QuerySet(object):
 
         return mongo_query
 
+    def first(self):
+        """Retrieve the first object matching the query.
+        """
+        result = self._collection.find_one(self._query)
+        if result is not None:
+            result = self._document._from_son(result)
+        return result
+
+    def with_id(self, object_id):
+        """Retrieve the object matching the _id provided.
+        """
+        if not isinstance(object_id, pymongo.objectid.ObjectId):
+            object_id = pymongo.objectid.ObjectId(object_id)
+
+        result = self._collection.find_one(object_id)
+        if result is not None:
+            result = self._document._from_son(result)
+        return result
+
     def next(self):
         """Wrap the result in a Document object.
         """
@@ -74,42 +97,22 @@ class QuerySet(object):
         return self
 
 
-class CollectionManager(object):
+class QuerySetManager(object):
-    
+
     def __init__(self, document):
-        """Set up the collection manager for a specific document.
-        """
         db = _get_db()
         self._document = document
         self._collection_name = document._meta['collection']
         # This will create the collection if it doesn't exist
         self._collection = db[self._collection_name]
 
-    def _save_document(self, document):
+    def __get__(self, instance, owner):
-        """Save the provided document to the collection.
+        """Descriptor for instantiating a new QuerySet object when 
+        Document.objects is accessed.
         """
-        _id = self._collection.save(document._to_mongo())
+        if instance is not None:
-        document._id = _id
+            # Document class being used rather than a document object
-
+            return self
-    def find(self, **query):
+        
-        """Query the collection for documents matching the provided query.
+        # self._document should be the same as owner
-        """
+        return QuerySet(self._document, self._collection)
-        return QuerySet(self._document, self._collection, query)
-
-    def find_one(self, object_id=None, **query):
-        """Query the collection for document matching the provided query.
-        """
-        if object_id:
-            # Use just object_id if provided
-            if not isinstance(object_id, pymongo.objectid.ObjectId):
-                object_id = pymongo.objectid.ObjectId(object_id)
-            query = object_id
-        else:
-            # Otherwise, use the query provided
-            query = QuerySet._transform_query(**query)
-            query['_types'] = self._document._class_name
-
-        result = self._collection.find_one(query)
-        if result is not None:
-            result = self._document._from_son(result)
-        return result

tests/document.py

@@ -104,13 +104,13 @@ class DocumentTest(unittest.TestCase):
         Human().save()
         Dog().save()
 
-        classes = [obj.__class__ for obj in Animal.objects.find()]
+        classes = [obj.__class__ for obj in Animal.objects]
         self.assertEqual(classes, [Animal, Fish, Mammal, Human, Dog])
 
-        classes = [obj.__class__ for obj in Mammal.objects.find()]
+        classes = [obj.__class__ for obj in Mammal.objects]
         self.assertEqual(classes, [Mammal, Human, Dog])
 
-        classes = [obj.__class__ for obj in Human.objects.find()]
+        classes = [obj.__class__ for obj in Human.objects]
         self.assertEqual(classes, [Human])
 
         Animal.drop_collection()
@@ -254,7 +254,7 @@ class DocumentTest(unittest.TestCase):
         post.author = author
         post.save()
 
-        post_obj = BlogPost.objects.find_one()
+        post_obj = BlogPost.objects.first()
 
         # Test laziness
         self.assertTrue(isinstance(post_obj._data['author'], 
@@ -266,7 +266,7 @@ class DocumentTest(unittest.TestCase):
         post_obj.author.age = 25
         post_obj.author.save()
 
-        author = self.Person.objects.find_one(name='Test User')
+        author = self.Person.objects(name='Test User').first()
         self.assertEqual(author.age, 25)
 
         BlogPost.drop_collection()

tests/queryset.py

@@ -1,11 +1,11 @@
 import unittest
 import pymongo
 
-from mongoengine.collection import CollectionManager, QuerySet
+from mongoengine.queryset import QuerySet
 from mongoengine import *
 
 
-class CollectionManagerTest(unittest.TestCase):
+class QuerySetTest(unittest.TestCase):
     
     def setUp(self):
         connect(db='mongoenginetest')
@@ -18,8 +18,8 @@ class CollectionManagerTest(unittest.TestCase):
     def test_initialisation(self):
         """Ensure that CollectionManager is correctly initialised.
         """
-        self.assertTrue(isinstance(self.Person.objects, CollectionManager))
+        self.assertTrue(isinstance(self.Person.objects, QuerySet))
-        self.assertEqual(self.Person.objects._collection_name, 
+        self.assertEqual(self.Person.objects._collection.name(), 
                          self.Person._meta['collection'])
         self.assertTrue(isinstance(self.Person.objects._collection,
                                    pymongo.collection.Collection))
@@ -45,7 +45,7 @@ class CollectionManagerTest(unittest.TestCase):
         person2.save()
 
         # Find all people in the collection
-        people = self.Person.objects.find()
+        people = self.Person.objects
         self.assertEqual(people.count(), 2)
         results = list(people)
         self.assertTrue(isinstance(results[0], self.Person))
@@ -57,19 +57,19 @@ class CollectionManagerTest(unittest.TestCase):
         self.assertEqual(results[1].age, 30)
 
         # Use a query to filter the people found to just person1
-        people = self.Person.objects.find(age=20)
+        people = self.Person.objects(age=20)
         self.assertEqual(people.count(), 1)
         person = people.next()
         self.assertEqual(person.name, "User A")
         self.assertEqual(person.age, 20)
 
         # Test limit
-        people = list(self.Person.objects.find().limit(1))
+        people = list(self.Person.objects.limit(1))
         self.assertEqual(len(people), 1)
         self.assertEqual(people[0].name, 'User A')
 
         # Test skip
-        people = list(self.Person.objects.find().skip(1))
+        people = list(self.Person.objects.skip(1))
         self.assertEqual(len(people), 1)
         self.assertEqual(people[0].name, 'User B')
 
@@ -82,20 +82,20 @@ class CollectionManagerTest(unittest.TestCase):
         person2.save()
 
         # Retrieve the first person from the database
-        person = self.Person.objects.find_one()
+        person = self.Person.objects.first()
         self.assertTrue(isinstance(person, self.Person))
         self.assertEqual(person.name, "User A")
         self.assertEqual(person.age, 20)
 
         # Use a query to filter the people found to just person2
-        person = self.Person.objects.find_one(age=30)
+        person = self.Person.objects(age=30).first()
         self.assertEqual(person.name, "User B")
 
-        person = self.Person.objects.find_one(age__lt=30)
+        person = self.Person.objects(age__lt=30).first()
         self.assertEqual(person.name, "User A")
         
         # Find a document using just the object id
-        person = self.Person.objects.find_one(person1._id)
+        person = self.Person.objects.with_id(person1._id)
         self.assertEqual(person.name, "User A")
 
     def test_find_embedded(self):
@@ -112,7 +112,7 @@ class CollectionManagerTest(unittest.TestCase):
         post.author = User(name='Test User')
         post.save()
 
-        result = BlogPost.objects.find_one()
+        result = BlogPost.objects.first()
         self.assertTrue(isinstance(result.author, User))
         self.assertEqual(result.author.name, 'Test User')
         
@@ -125,13 +125,13 @@ class CollectionManagerTest(unittest.TestCase):
         self.Person(name="User B", age=30).save()
         self.Person(name="User C", age=40).save()
 
-        self.assertEqual(self.Person.objects.find().count(), 3)
+        self.assertEqual(self.Person.objects.count(), 3)
 
-        self.Person.objects.find(age__lt=30).delete()
+        self.Person.objects(age__lt=30).delete()
-        self.assertEqual(self.Person.objects.find().count(), 2)
+        self.assertEqual(self.Person.objects.count(), 2)
 
-        self.Person.objects.find().delete()
+        self.Person.objects.delete()
-        self.assertEqual(self.Person.objects.find().count(), 0)
+        self.assertEqual(self.Person.objects.count(), 0)
 
     def tearDown(self):
         self.Person.drop_collection()