Added pretty query format, cursor limit and skip

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))

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.
         """