diff --git a/mongoengine/queryset.py b/mongoengine/queryset.py index 73a45299..016430d1 100644 --- a/mongoengine/queryset.py +++ b/mongoengine/queryset.py @@ -916,20 +916,27 @@ class QuerySet(object): """ return self.exec_js(average_func, field) - def item_frequencies(self, list_field, normalize=False): - """Returns a dictionary of all items present in a list field across + def item_frequencies(self, field, normalize=False): + """Returns a dictionary of all items present in a field across the whole queried set of documents, and their corresponding frequency. This is useful for generating tag clouds, or searching documents. - :param list_field: the list field to use + If the field is a :class:`~mongoengine.ListField`, the items within + each list will be counted individually. + + :param field: the field to use :param normalize: normalize the results so they add to 1.0 """ freq_func = """ - function(listField) { + function(field) { if (options.normalize) { var total = 0.0; db[collection].find(query).forEach(function(doc) { - total += doc[listField].length; + if (doc[field].constructor == Array) { + total += doc[field].length; + } else { + total++; + } }); } @@ -939,14 +946,19 @@ class QuerySet(object): inc /= total; } db[collection].find(query).forEach(function(doc) { - doc[listField].forEach(function(item) { + if (doc[field].constructor == Array) { + doc[field].forEach(function(item) { + frequencies[item] = inc + (frequencies[item] || 0); + }); + } else { + var item = doc[field]; frequencies[item] = inc + (frequencies[item] || 0); - }); + } }); return frequencies; } """ - return self.exec_js(freq_func, list_field, normalize=normalize) + return self.exec_js(freq_func, field, normalize=normalize) def __repr__(self): limit = REPR_OUTPUT_SIZE + 1 diff --git a/tests/queryset.py b/tests/queryset.py index ab28ff3d..62cf4954 100644 --- a/tests/queryset.py +++ b/tests/queryset.py @@ -973,7 +973,7 @@ class QuerySetTest(unittest.TestCase): BlogPost(hits=1, tags=['music', 'film', 'actors']).save() BlogPost(hits=2, tags=['music']).save() - BlogPost(hits=3, tags=['music', 'actors']).save() + BlogPost(hits=2, tags=['music', 'actors']).save() f = BlogPost.objects.item_frequencies('tags') f = dict((key, int(val)) for key, val in f.items()) @@ -995,6 +995,13 @@ class QuerySetTest(unittest.TestCase): self.assertAlmostEqual(f['actors'], 2.0/6.0) self.assertAlmostEqual(f['film'], 1.0/6.0) + # Check item_frequencies works for non-list fields + f = BlogPost.objects.item_frequencies('hits') + f = dict((key, int(val)) for key, val in f.items()) + self.assertEqual(set(['1', '2']), set(f.keys())) + self.assertEqual(f['1'], 1) + self.assertEqual(f['2'], 2) + BlogPost.drop_collection() def test_average(self):