From d4fc5c926093affc7ea487310e7f11157219e0e4 Mon Sep 17 00:00:00 2001 From: Harry Marr Date: Sun, 22 Nov 2009 16:46:08 +0000 Subject: [PATCH] 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 | 4 +- mongoengine/document.py | 3 +- mongoengine/fields.py | 3 +- mongoengine/{collection.py => queryset.py} | 73 +++++++++++----------- tests/document.py | 10 +-- tests/{collection.py => queryset.py} | 36 +++++------ 6 files changed, 67 insertions(+), 62 deletions(-) rename mongoengine/{collection.py => queryset.py} (65%) rename tests/{collection.py => queryset.py} (80%) diff --git a/mongoengine/base.py b/mongoengine/base.py index b180231a..e5d8d002 100644 --- a/mongoengine/base.py +++ b/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 diff --git a/mongoengine/document.py b/mongoengine/document.py index 315776aa..ab6b74a8 100644 --- a/mongoengine/document.py +++ b/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): diff --git a/mongoengine/fields.py b/mongoengine/fields.py index 93b21f1e..70e58be8 100644 --- a/mongoengine/fields.py +++ b/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) diff --git a/mongoengine/collection.py b/mongoengine/queryset.py similarity index 65% rename from mongoengine/collection.py rename to mongoengine/queryset.py index 34329790..d5938795 100644 --- a/mongoengine/collection.py +++ b/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 = QuerySet._transform_query(**query) - self._query['_types'] = self._document._class_name + 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): - """Save the provided document to the collection. + def __get__(self, instance, owner): + """Descriptor for instantiating a new QuerySet object when + Document.objects is accessed. """ - _id = self._collection.save(document._to_mongo()) - document._id = _id - - def find(self, **query): - """Query the collection for documents matching the provided query. - """ - 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 + if instance is not None: + # Document class being used rather than a document object + return self + + # self._document should be the same as owner + return QuerySet(self._document, self._collection) diff --git a/tests/document.py b/tests/document.py index 2f2fef3c..4842934c 100644 --- a/tests/document.py +++ b/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() diff --git a/tests/collection.py b/tests/queryset.py similarity index 80% rename from tests/collection.py rename to tests/queryset.py index d77f304c..2c650b4c 100644 --- a/tests/collection.py +++ b/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.assertEqual(self.Person.objects._collection_name, + self.assertTrue(isinstance(self.Person.objects, QuerySet)) + 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.assertEqual(self.Person.objects.find().count(), 2) + self.Person.objects(age__lt=30).delete() + self.assertEqual(self.Person.objects.count(), 2) - self.Person.objects.find().delete() - self.assertEqual(self.Person.objects.find().count(), 0) + self.Person.objects.delete() + self.assertEqual(self.Person.objects.count(), 0) def tearDown(self): self.Person.drop_collection()