extending support for queryset.sum and queryset.average methods

mongoengine/queryset/queryset.py

@@ -1066,11 +1066,27 @@ class QuerySet(object):
         .. versionchanged:: 0.5 - updated to map_reduce as db.eval doesnt work
             with sharding.
         """
-        map_func = Code("""
+        map_func = """
             function() {
-                emit(1, this[field] || 0);
+                var path = '{{~%(field)s}}'.split('.'),
+                field = this;
+
+                for (p in path) {
+                    if (typeof field != 'undefined')
+                       field = field[path[p]];
+                    else
+                       break;
+                }
+        
+                if (field && field.constructor == Array) {
+                    field.forEach(function(item) {
+                        emit(1, item||0);
+                    });
+                } else if (typeof field != 'undefined') {
+                    emit(1, field||0);
+                }
             }
-        """, scope={'field': field})
+        """ % dict(field=field)
 
         reduce_func = Code("""
             function(key, values) {
@@ -1096,13 +1112,28 @@ class QuerySet(object):
         .. versionchanged:: 0.5 - updated to map_reduce as db.eval doesnt work
             with sharding.
         """
-        map_func = Code("""
+        map_func = """
             function() {
-                if (this.hasOwnProperty(field))
+                var path = '{{~%(field)s}}'.split('.'),
-                    emit(1, {t: this[field] || 0, c: 1});
+                field = this;
-            }
-        """, scope={'field': field})
 
+                for (p in path) {
+                    if (typeof field != 'undefined')
+                       field = field[path[p]];
+                    else
+                       break;
+                }
+        
+                if (field && field.constructor == Array) {
+                    field.forEach(function(item) {
+                        emit(1, {t: item||0, c: 1});
+                    });
+                } else if (typeof field != 'undefined') {
+                    emit(1, {t: field||0, c: 1});
+                }
+            }
+        """ % dict(field=field)
+    
         reduce_func = Code("""
             function(key, values) {
                 var out = {t: 0, c: 0};
@@ -1263,8 +1294,8 @@ class QuerySet(object):
     def _item_frequencies_map_reduce(self, field, normalize=False):
         map_func = """
             function() {
-                var path = '{{~%(field)s}}'.split('.');
+                var path = '{{~%(field)s}}'.split('.'),
-                var field = this;
+                field = this;
 
                 for (p in path) {
                     if (typeof field != 'undefined')

tests/queryset/queryset.py

@@ -2208,6 +2208,75 @@ class QuerySetTest(unittest.TestCase):
         self.Person(name='ageless person').save()
         self.assertEqual(int(self.Person.objects.average('age')), avg)
 
+    def test_embedded_average(self):
+        class Pay(EmbeddedDocument):
+            value = DecimalField()
+
+        class Doc(Document):
+            name = StringField()
+            pay = EmbeddedDocumentField(
+                Pay)
+
+        Doc.drop_collection()
+
+        Doc(name=u"Wilson Junior",
+            pay=Pay(value=150)).save()
+
+        Doc(name=u"Isabella Luanna",
+            pay=Pay(value=530)).save()
+
+        Doc(name=u"Tayza mariana",
+            pay=Pay(value=165)).save()
+
+        Doc(name=u"Eliana Costa",
+            pay=Pay(value=115)).save()
+
+        self.assertEqual(
+            Doc.objects.average('pay.value'),
+            240)
+
+    def test_embedded_array_average(self):
+        class Pay(EmbeddedDocument):
+            values = ListField(DecimalField())
+
+        class Doc(Document):
+            name = StringField()
+            pay = EmbeddedDocumentField(
+                Pay)
+
+        Doc.drop_collection()
+
+        Doc(name=u"Wilson Junior",
+            pay=Pay(values=[150, 100])).save()
+
+        Doc(name=u"Isabella Luanna",
+            pay=Pay(values=[530, 100])).save()
+
+        Doc(name=u"Tayza mariana",
+            pay=Pay(values=[165, 100])).save()
+
+        Doc(name=u"Eliana Costa",
+            pay=Pay(values=[115, 100])).save()
+
+        self.assertEqual(
+            Doc.objects.average('pay.values'),
+            170)
+
+    def test_array_average(self):
+        class Doc(Document):
+            values = ListField(DecimalField())
+
+        Doc.drop_collection()
+
+        Doc(values=[150, 100]).save()
+        Doc(values=[530, 100]).save()
+        Doc(values=[165, 100]).save()
+        Doc(values=[115, 100]).save()
+
+        self.assertEqual(
+            Doc.objects.average('values'),
+            170)
+
     def test_sum(self):
         """Ensure that field can be summed over correctly.
         """
@@ -2220,6 +2289,77 @@ class QuerySetTest(unittest.TestCase):
         self.Person(name='ageless person').save()
         self.assertEqual(int(self.Person.objects.sum('age')), sum(ages))
 
+    def test_embedded_sum(self):
+        class Pay(EmbeddedDocument):
+            value = DecimalField()
+
+        class Doc(Document):
+            name = StringField()
+            pay = EmbeddedDocumentField(
+                Pay)
+
+        Doc.drop_collection()
+
+        Doc(name=u"Wilson Junior",
+            pay=Pay(value=150)).save()
+
+        Doc(name=u"Isabella Luanna",
+            pay=Pay(value=530)).save()
+
+        Doc(name=u"Tayza mariana",
+            pay=Pay(value=165)).save()
+
+        Doc(name=u"Eliana Costa",
+            pay=Pay(value=115)).save()
+
+        self.assertEqual(
+            Doc.objects.sum('pay.value'),
+            960)
+
+
+    def test_embedded_array_sum(self):
+        class Pay(EmbeddedDocument):
+            values = ListField(DecimalField())
+
+        class Doc(Document):
+            name = StringField()
+            pay = EmbeddedDocumentField(
+                Pay)
+
+        Doc.drop_collection()
+
+        Doc(name=u"Wilson Junior",
+            pay=Pay(values=[150, 100])).save()
+
+        Doc(name=u"Isabella Luanna",
+            pay=Pay(values=[530, 100])).save()
+
+        Doc(name=u"Tayza mariana",
+            pay=Pay(values=[165, 100])).save()
+
+        Doc(name=u"Eliana Costa",
+            pay=Pay(values=[115, 100])).save()
+
+        self.assertEqual(
+            Doc.objects.sum('pay.values'),
+            1360)
+
+    def test_array_sum(self):
+        class Doc(Document):
+            values = ListField(DecimalField())
+
+        Doc.drop_collection()
+
+        Doc(values=[150, 100]).save()
+        Doc(values=[530, 100]).save()
+        Doc(values=[165, 100]).save()
+        Doc(values=[115, 100]).save()
+
+        self.assertEqual(
+            Doc.objects.sum('values'),
+            1360)
+
+            
     def test_distinct(self):
         """Ensure that the QuerySet.distinct method works.
         """