From 8ec6fecd2394cf47fb29b85510f577ea95495709 Mon Sep 17 00:00:00 2001 From: Harry Marr Date: Thu, 19 Nov 2009 01:09:58 +0000 Subject: [PATCH] Added basic querying - find and find_one --- mongomap/base.py | 12 +++++-- mongomap/collection.py | 37 +++++++++++++++++++- mongomap/document.py | 5 ++- mongomap/fields.py | 3 ++ tests/collection.py | 77 +++++++++++++++++++++++++++++++++++++----- tests/document.py | 1 + tests/fields.py | 18 ++++++++++ 7 files changed, 141 insertions(+), 12 deletions(-) diff --git a/mongomap/base.py b/mongomap/base.py index c2bfa0fe..5401ea2f 100644 --- a/mongomap/base.py +++ b/mongomap/base.py @@ -44,7 +44,7 @@ class BaseField(object): try: value = self._to_python(value) self._validate(value) - except ValueError: + except (ValueError, AttributeError): raise ValidationError('Invalid value for field of type "' + self.__class__.__name__ + '"') elif self.required: @@ -145,7 +145,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) - setattr(new_class, 'collection', CollectionManager(new_class)) + new_class.objects = CollectionManager(new_class) return new_class @@ -204,3 +204,11 @@ class BaseDocument(object): if value is not None: data[field_name] = field._to_mongo(value) return data + + @classmethod + def _from_son(cls, son): + """Create an instance of a Document (subclass) from a PyMongo SOM. + """ + data = dict((str(key), value) for key, value in son.items()) + return cls(**data) + diff --git a/mongomap/collection.py b/mongomap/collection.py index 5016a1a5..a3debdad 100644 --- a/mongomap/collection.py +++ b/mongomap/collection.py @@ -1,5 +1,29 @@ from connection import _get_db + +class QuerySet(object): + """A set of results returned from a query. Wraps a MongoDB cursor, + providing Document objects as the results. + """ + + def __init__(self, document, cursor): + self._document = document + self._cursor = cursor + + def next(self): + """Wrap the result in a Document object. + """ + return self._document._from_son(self._cursor.next()) + + def count(self): + """Count the selected elements in the query. + """ + return self._cursor.count() + + def __iter__(self): + return self + + class CollectionManager(object): def __init__(self, document): @@ -14,4 +38,15 @@ class CollectionManager(object): def _save_document(self, document): """Save the provided document to the collection. """ - _id = self._collection.save(document) + _id = self._collection.save(document._to_mongo()) + document._id = _id + + def find(self, query=None): + """Query the collection for document matching the provided query. + """ + return QuerySet(self._document, self._collection.find(query)) + + def find_one(self, query=None): + """Query the collection for document matching the provided query. + """ + return self._document._from_son(self._collection.find_one(query)) diff --git a/mongomap/document.py b/mongomap/document.py index 7b7f0289..e83dd452 100644 --- a/mongomap/document.py +++ b/mongomap/document.py @@ -14,4 +14,7 @@ class Document(BaseDocument): __metaclass__ = TopLevelDocumentMetaclass def save(self): - self.collection._save_document(self._to_mongo()) + """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) diff --git a/mongomap/fields.py b/mongomap/fields.py index 3003e3c9..dd4cbebe 100644 --- a/mongomap/fields.py +++ b/mongomap/fields.py @@ -59,6 +59,8 @@ class EmbeddedDocumentField(BaseField): super(EmbeddedDocumentField, self).__init__(**kwargs) def _to_python(self, value): + if not isinstance(value, self.document): + return self.document._from_son(value) return value def _to_mongo(self, value): @@ -68,6 +70,7 @@ class EmbeddedDocumentField(BaseField): """Make sure that the document instance is an instance of the EmbeddedDocument subclass provided when the document was defined. """ + # Using isinstance also works for subclasses of self.document if not isinstance(value, self.document): raise ValidationError('Invalid embedded document instance ' 'provided to an EmbeddedDocumentField') diff --git a/tests/collection.py b/tests/collection.py index 4f2bc71d..b1e3f933 100644 --- a/tests/collection.py +++ b/tests/collection.py @@ -2,6 +2,7 @@ import unittest import pymongo from mongomap.collection import CollectionManager +from mongomap.connection import _get_db from mongomap import * @@ -15,19 +16,79 @@ class CollectionManagerTest(unittest.TestCase): age = IntField() self.Person = Person + self.db = _get_db() + self.db.drop_collection(self.Person._meta['collection']) + def test_initialisation(self): """Ensure that CollectionManager is correctly initialised. """ - class Person(Document): - name = StringField() - age = IntField() - - self.assertTrue(isinstance(Person.collection, CollectionManager)) - self.assertEqual(Person.collection._collection_name, - Person._meta['collection']) - self.assertTrue(isinstance(Person.collection._collection, + self.assertTrue(isinstance(self.Person.objects, CollectionManager)) + self.assertEqual(self.Person.objects._collection_name, + self.Person._meta['collection']) + self.assertTrue(isinstance(self.Person.objects._collection, pymongo.collection.Collection)) + def test_find(self): + """Ensure that a query returns a valid set of results. + """ + person1 = self.Person(name="User A", age=20) + person1.save() + person2 = self.Person(name="User B", age=30) + person2.save() + + # Find all people in the collection + people = self.Person.objects.find() + self.assertEqual(people.count(), 2) + results = list(people) + self.assertTrue(isinstance(results[0], self.Person)) + self.assertEqual(results[0].name, "User A") + self.assertEqual(results[0].age, 20) + self.assertEqual(results[1].name, "User B") + self.assertEqual(results[1].age, 30) + + # Use a query to filter the people found to just person1 + people = self.Person.objects.find({'age': 20}) + self.assertEqual(people.count(), 1) + person = people.next() + self.assertEqual(person.name, "User A") + self.assertEqual(person.age, 20) + + def test_find_one(self): + """Ensure that a query using find_one returns a valid result. + """ + person1 = self.Person(name="User A", age=20) + person1.save() + person2 = self.Person(name="User B", age=30) + person2.save() + + # Retrieve the first person from the database + person = self.Person.objects.find_one() + 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}) + self.assertEqual(person.name, "User B") + + def test_find_embedded(self): + """Ensure that an embedded document is properly returned from a query. + """ + class User(EmbeddedDocument): + name = StringField() + + class BlogPost(Document): + content = StringField() + author = EmbeddedDocumentField(User) + + post = BlogPost(content='Had a good coffee today...') + post.author = User(name='Test User') + post.save() + + result = BlogPost.objects.find_one() + self.assertTrue(isinstance(result.author, User)) + self.assertEqual(result.author.name, 'Test User') + if __name__ == '__main__': unittest.main() diff --git a/tests/document.py b/tests/document.py index 7a6e1c4f..177b5585 100644 --- a/tests/document.py +++ b/tests/document.py @@ -97,6 +97,7 @@ class DocumentTest(unittest.TestCase): person_obj = collection.find_one({'name': 'Test User'}) self.assertEqual(person_obj['name'], 'Test User') self.assertEqual(person_obj['age'], 30) + self.assertEqual(str(person_obj['_id']), person._id) def test_save_custom_id(self): """Ensure that a document may be saved with a custom _id. diff --git a/tests/fields.py b/tests/fields.py index c4af2be0..e66974c8 100644 --- a/tests/fields.py +++ b/tests/fields.py @@ -103,6 +103,24 @@ class FieldTest(unittest.TestCase): person.preferences = PersonPreferences(food='Cheese', number=47) self.assertEqual(person.preferences.food, 'Cheese') + def test_embedded_document_inheritance(self): + """Ensure that subclasses of embedded documents may be provided to + EmbeddedDocumentFields of the superclass' type. + """ + class User(EmbeddedDocument): + name = StringField() + + class PowerUser(User): + power = IntField() + + class BlogPost(Document): + content = StringField() + author = EmbeddedDocumentField(User) + + post = BlogPost(content='What I did today...') + post.author = User(name='Test User') + post.author = PowerUser(name='Test User', power=47) + if __name__ == '__main__': unittest.main()