From 6438bd52b7cf607cd5b2d01e2b98667b46116de0 Mon Sep 17 00:00:00 2001
From: Harry Marr <harry.marr@gmail.com>
Date: Sun, 20 Dec 2009 17:17:56 +0000
Subject: [PATCH] Added item_frequencies to QuerySet

---
 mongoengine/queryset.py | 20 ++++++++++++++++++++
 tests/queryset.py       | 29 +++++++++++++++++++++++++++++
 2 files changed, 49 insertions(+)

diff --git a/mongoengine/queryset.py b/mongoengine/queryset.py
index 034b8476..68c27067 100644
--- a/mongoengine/queryset.py
+++ b/mongoengine/queryset.py
@@ -159,6 +159,26 @@ class QuerySet(object):
     def __iter__(self):
         return self
 
+    def item_frequencies(self, list_field):
+        """Returns a dictionary of all items present in a list field across
+        the whole queried set of documents, and their corresponding frequency.
+        This is useful for generating tag clouds, or searching documents. 
+        """
+        freq_func = """
+            function(collection, query, listField) {
+                var frequencies = {};
+                db[collection].find(query).forEach(function(doc) {
+                    doc[listField].forEach(function(item) {
+                        frequencies[item] = 1 + (frequencies[item] || 0);
+                    });
+                });
+                return frequencies;
+            }
+        """
+        db = _get_db()
+        collection = self._document._meta['collection']
+        return db.eval(freq_func, collection, self._query, list_field)
+
 
 class QuerySetManager(object):
 
diff --git a/tests/queryset.py b/tests/queryset.py
index b73bab17..461ad34b 100644
--- a/tests/queryset.py
+++ b/tests/queryset.py
@@ -185,6 +185,35 @@ class QuerySetTest(unittest.TestCase):
         ages = [p.age for p in self.Person.objects.order_by('-name')]
         self.assertEqual(ages, [30, 40, 20])
 
+    def test_item_frequencies(self):
+        """Ensure that item frequencies are properly generated from lists.
+        """
+        class BlogPost(Document):
+            hits = IntField()
+            tags = ListField(StringField())
+
+        BlogPost.drop_collection()
+
+        BlogPost(hits=1, tags=['music', 'film', 'actors']).save()
+        BlogPost(hits=2, tags=['music']).save()
+        BlogPost(hits=3, tags=['music', 'actors']).save()
+
+        f = BlogPost.objects.item_frequencies('tags')
+        f = dict((key, int(val)) for key, val in f.items())
+        self.assertEqual(set(['music', 'film', 'actors']), set(f.keys()))
+        self.assertEqual(f['music'], 3)
+        self.assertEqual(f['actors'], 2)
+        self.assertEqual(f['film'], 1)
+
+        # Ensure query is taken into account
+        f = BlogPost.objects(hits__gt=1).item_frequencies('tags')
+        f = dict((key, int(val)) for key, val in f.items())
+        self.assertEqual(set(['music', 'actors']), set(f.keys()))
+        self.assertEqual(f['music'], 2)
+        self.assertEqual(f['actors'], 1)
+
+        BlogPost.drop_collection()
+
     def tearDown(self):
         self.Person.drop_collection()