diff --git a/mongomap/collection.py b/mongomap/collection.py index a3debdad..db4c0760 100644 --- a/mongomap/collection.py +++ b/mongomap/collection.py @@ -20,6 +20,19 @@ class QuerySet(object): """ return self._cursor.count() + def limit(self, n): + """Limit the number of returned documents to. + """ + self._cursor.limit(n) + # Return self to allow chaining + return self + + def skip(self, n): + """Skip n documents before returning the results. + """ + self._cursor.skip(n) + return self + def __iter__(self): return self @@ -41,12 +54,35 @@ class CollectionManager(object): _id = self._collection.save(document._to_mongo()) document._id = _id - def find(self, query=None): + def _transform_query(self, **query): + """Transform a query from Django-style format to Mongo format. + """ + operators = ['neq', 'gt', 'gte', 'lt', 'lte', 'in', 'nin', 'mod', + 'all', 'size', 'exists'] + + mongo_query = {} + for key, value in query.items(): + parts = key.split('__') + # Check for an operator and transform to mongo-style if there is + if parts[-1] in operators: + op = parts.pop() + value = {'$' + op: value} + + key = '.'.join(parts) + mongo_query[key] = value + + return mongo_query + + def find(self, **query): """Query the collection for document matching the provided query. """ + if query: + query = self._transform_query(**query) return QuerySet(self._document, self._collection.find(query)) - def find_one(self, query=None): + def find_one(self, **query): """Query the collection for document matching the provided query. """ + if query: + query = self._transform_query(**query) return self._document._from_son(self._collection.find_one(query)) diff --git a/tests/collection.py b/tests/collection.py index b1e3f933..4de6cbbc 100644 --- a/tests/collection.py +++ b/tests/collection.py @@ -28,6 +28,19 @@ class CollectionManagerTest(unittest.TestCase): self.assertTrue(isinstance(self.Person.objects._collection, pymongo.collection.Collection)) + def test_transform_query(self): + """Ensure that the _transform_query function operates correctly. + """ + manager = self.Person().objects + self.assertEqual(manager._transform_query(name='test', age=30), + {'name': 'test', 'age': 30}) + self.assertEqual(manager._transform_query(age__lt=30), + {'age': {'$lt': 30}}) + self.assertEqual(manager._transform_query(friend__age__gte=30), + {'friend.age': {'$gte': 30}}) + self.assertEqual(manager._transform_query(name__exists=True), + {'name': {'$exists': True}}) + def test_find(self): """Ensure that a query returns a valid set of results. """ @@ -47,12 +60,22 @@ 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.find(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)) + self.assertEqual(len(people), 1) + self.assertEqual(people[0].name, 'User A') + + # Test skip + people = list(self.Person.objects.find().skip(1)) + self.assertEqual(len(people), 1) + self.assertEqual(people[0].name, 'User B') + def test_find_one(self): """Ensure that a query using find_one returns a valid result. """ @@ -68,9 +91,12 @@ class CollectionManagerTest(unittest.TestCase): 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.find_one(age=30) self.assertEqual(person.name, "User B") + person = self.Person.objects.find_one(age__lt=30) + self.assertEqual(person.name, "User A") + def test_find_embedded(self): """Ensure that an embedded document is properly returned from a query. """