Merge pull request #703 from wpjunior/aggregate-framework

Simple aggregation framework wrapper

mongoengine/queryset/base.py

@@ -974,8 +974,33 @@ class BaseQuerySet(object):
         son_data = json_util.loads(json_data)
         return [self._document._from_son(data) for data in son_data]
 
-    # JS functionality
+    def aggregate(self, *pipeline, **kwargs):
+        """
+        Perform a aggreggate function based in your queryset params
+        :param pipeline: list of agreggation commands,
+            see: http://docs.mongodb.org/manual/core/aggregation-pipeline/
 
+        .. versionadded:: 0.9
+        """
+        initial_pipeline = []
+
+        if self._query:
+            initial_pipeline.append({'$match': self._query})
+
+        if self._ordering:
+            initial_pipeline.append({'$sort': dict(self._ordering)})
+
+        if self._limit is not None:
+            initial_pipeline.append({'$limit': self._limit})
+
+        if self._skip is not None:
+            initial_pipeline.append({'$skip': self._skip})
+
+        pipeline = initial_pipeline + list(pipeline)
+
+        return self._collection.aggregate(pipeline, cursor={}, **kwargs)
+
+    # JS functionality
     def map_reduce(self, map_f, reduce_f, output, finalize_f=None, limit=None,
                    scope=None):
         """Perform a map/reduce query using the current query spec

tests/queryset/queryset.py

@@ -4341,6 +4341,55 @@ class QuerySetTest(unittest.TestCase):
             self.assertTrue(Person.objects._has_data(),
                             'Cursor has data and returned False')
 
+    def test_queryset_aggregation_framework(self):
+        class Person(Document):
+            name = StringField()
+            age = IntField()
+
+        Person.drop_collection()
+
+        p1 = Person(name="Isabella Luanna", age=16)
+        p1.save()
+
+        p2 = Person(name="Wilson Junior", age=21)
+        p2.save()
+
+        p3 = Person(name="Sandra Mara", age=37)
+        p3.save()
+
+        data = Person.objects(age__lte=22).aggregate(
+            {'$project': {'name': {'$toUpper': '$name'}}}
+        )
+
+        self.assertEqual(list(data), [
+            {'_id': p1.pk, 'name': "ISABELLA LUANNA"},
+            {'_id': p2.pk, 'name': "WILSON JUNIOR"}
+        ])
+
+        data = Person.objects(age__lte=22).order_by('-name').aggregate(
+            {'$project': {'name': {'$toUpper': '$name'}}}
+        )
+
+        self.assertEqual(list(data), [
+            {'_id': p2.pk, 'name': "WILSON JUNIOR"},
+            {'_id': p1.pk, 'name': "ISABELLA LUANNA"}
+        ])
+
+        data = Person.objects(
+            age__gte=17, age__lte=40).order_by('-age').aggregate(
+                {'$group': {
+                    '_id': None,
+                    'total': {'$sum': 1},
+                    'avg': {'$avg': '$age'}
+                }
+                }
+
+        )
+
+        self.assertEqual(list(data), [
+            {'_id': None, 'avg': 29, 'total': 2}
+        ])
+
 
 if __name__ == '__main__':
     unittest.main()