From 60c8254f58986f43a25f9190bb97fb07b0742b64 Mon Sep 17 00:00:00 2001 From: Ross Lawley Date: Wed, 25 May 2011 11:10:42 +0100 Subject: [PATCH] Tweaks to item_frequencies Updated to use a ternary statement and added tests Refs #124 #122 Thanks to @nickvlku for the code. --- mongoengine/queryset.py | 12 ++---------- tests/queryset.py | 13 ++++++++----- 2 files changed, 10 insertions(+), 15 deletions(-) diff --git a/mongoengine/queryset.py b/mongoengine/queryset.py index 54d4845e..f5020ab8 100644 --- a/mongoengine/queryset.py +++ b/mongoengine/queryset.py @@ -1298,19 +1298,11 @@ class QuerySet(object): db[collection].find(query).forEach(function(doc) { if (doc[field].constructor == Array) { doc[field].forEach(function(item) { - var preValue = 0; - if (!isNaN(frequencies[item])) { - preValue = frequencies[item]; - } - frequencies[item] = inc + preValue; + frequencies[item] = inc + (isNaN(frequencies[item]) ? 0: frequencies[item]); }); } else { var item = doc[field]; - var preValue = 0; - if (!isNaN(frequencies[item])) { - preValue = frequencies[item]; - } - frequencies[item] = inc + preValue; + frequencies[item] = inc + (isNaN(frequencies[item]) ? 0: frequencies[item]); } }); return frequencies; diff --git a/tests/queryset.py b/tests/queryset.py index 5cf08957..9b711406 100644 --- a/tests/queryset.py +++ b/tests/queryset.py @@ -1339,7 +1339,7 @@ class QuerySetTest(unittest.TestCase): BlogPost.drop_collection() BlogPost(hits=1, tags=['music', 'film', 'actors', 'watch']).save() - BlogPost(hits=2, tags=['music']).save() + BlogPost(hits=2, tags=['music', 'watch']).save() BlogPost(hits=2, tags=['music', 'actors']).save() f = BlogPost.objects.item_frequencies('tags') @@ -1347,20 +1347,23 @@ class QuerySetTest(unittest.TestCase): self.assertEqual(set(['music', 'film', 'actors', 'watch']), set(f.keys())) self.assertEqual(f['music'], 3) self.assertEqual(f['actors'], 2) + self.assertEqual(f['watch'], 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(set(['music', 'actors', 'watch']), set(f.keys())) self.assertEqual(f['music'], 2) self.assertEqual(f['actors'], 1) + self.assertEqual(f['watch'], 1) # Check that normalization works f = BlogPost.objects.item_frequencies('tags', normalize=True) - self.assertAlmostEqual(f['music'], 3.0/7.0) - self.assertAlmostEqual(f['actors'], 2.0/7.0) - self.assertAlmostEqual(f['film'], 1.0/7.0) + self.assertAlmostEqual(f['music'], 3.0/8.0) + self.assertAlmostEqual(f['actors'], 2.0/8.0) + self.assertAlmostEqual(f['watch'], 2.0/8.0) + self.assertAlmostEqual(f['film'], 1.0/8.0) # Check item_frequencies works for non-list fields f = BlogPost.objects.item_frequencies('hits')