diff --git a/pymongo/collection.py b/pymongo/collection.py index 6afcb484b1..e98e618afd 100644 --- a/pymongo/collection.py +++ b/pymongo/collection.py @@ -553,7 +553,8 @@ def options(self): # # Waiting on this because group command support for CodeWScope # wasn't added until 1.1 - def group(self, keys, condition, initial, reduce, command=False): + def group(self, keys, condition, initial, reduce, finalize=None, + command=False): """Perform a query similar to an SQL group by operation. Returns an array of grouped items. @@ -564,19 +565,33 @@ def group(self, keys, condition, initial, reduce, command=False): query specification) - `initial`: initial value of the aggregation counter object - `reduce`: aggregation function as a JavaScript string + - `finalize`: function to be called on each object in output list. - `command` (optional): if True, run the group as a command instead of in an eval - it is likely that this option will eventually be - deprecated and all groups will be run as commands + deprecated and all groups will be run as commands. Please only use + as a keyword argument, not as a positional argument. """ + + #for now support people passing command in its old position + if finalize in (True, False): + command = finalize + finalize = None + warnings.warn("Please only pass 'command' as a keyword argument." + ,DeprecationWarning) + if command: if not isinstance(reduce, Code): reduce = Code(reduce) - return self.__database._command({"group": - {"ns": self.__collection_name, - "$reduce": reduce, - "key": self._fields_list_to_dict(keys), - "cond": condition, - "initial": initial}})["retval"] + group = {"ns": self.__collection_name, + "$reduce": reduce, + "key": self._fields_list_to_dict(keys), + "cond": condition, + "initial": initial} + if finalize is not None: + if not isinstance(finalize, Code): + finalize = Code(finalize) + group["finalize"] = finalize + return self.__database._command({"group":group})["retval"] scope = {} if isinstance(reduce, Code): @@ -590,6 +605,7 @@ def group(self, keys, condition, initial, reduce, command=False): var c = db[ns].find(condition); var map = new Map(); var reduce_function = %s; + var finalize_function = %s; //function or null while (c.hasNext()) { var obj = c.next(); @@ -607,8 +623,18 @@ def group(self, keys, condition, initial, reduce, command=False): } reduce_function(obj, aggObj); } - return {"result": map.values()}; -}""" % reduce + + out = map.values(); + if (finalize_function !== null){ + for (var i=0; i < out.length; i++){ + var ret = finalize_function(out[i]); + if (ret !== undefined) + out[i] = ret; + } + } + + return {"result": out}; +}""" % (reduce, (finalize or 'null')); return self.__database.eval(Code(group_function, scope))["result"] def rename(self, new_name): diff --git a/test/test_collection.py b/test/test_collection.py index 91ca4051ec..7e22eb0e77 100644 --- a/test/test_collection.py +++ b/test/test_collection.py @@ -519,55 +519,68 @@ def test_group(self): db = self.db db.drop_collection("test") - self.assertEqual([], db.test.group([], {}, - {"count": 0}, - "function (obj, prev) { " - "prev.count++; }")) - self.assertEqual([], db.test.group([], {}, - {"count": 0}, - "function (obj, prev) { " - "prev.count++; }", command=True)) + def group_checker(args, expected): + eval = db.test.group(*args) + cmd = db.test.group(*args, command=True) + self.assertEqual(eval, expected) + self.assertEqual(cmd, expected) + self.assertEqual(eval, cmd) + #last one is not strictly necessary but there for completeness + + + args = [[], {}, + {"count": 0}, + "function (obj, prev) { prev.count++; }"] + expected = [] + group_checker(args, expected) db.test.save({"a": 2}) db.test.save({"b": 5}) db.test.save({"a": 1}) - self.assertEqual(3, db.test.group([], {}, - {"count": 0}, - "function (obj, prev) { " - "prev.count++; }")[0]["count"]) - self.assertEqual(3, db.test.group([], {}, - {"count": 0}, - "function (obj, prev) { " - "prev.count++; }", - command=True)[0]["count"]) - self.assertEqual(1, db.test.group([], - {"a": {"$gt": 1}}, - {"count": 0}, - "function (obj, prev) { " - "prev.count++; }")[0]["count"]) - self.assertEqual(1, db.test.group([], - {"a": {"$gt": 1}}, - {"count": 0}, - "function (obj, prev) { " - "prev.count++; }", - command=True)[0]["count"]) + args = [[], {}, + {"count": 0}, + "function (obj, prev) { prev.count++; }"] + expected = [{'count': 3}] + group_checker(args, expected) + + args = [[], + {"a": {"$gt": 1}}, + {"count": 0}, + "function (obj, prev) { prev.count++; }"] + expected = [{'count': 1}] + group_checker(args, expected) db.test.save({"a": 2, "b": 3}) + args = [["a"], {}, + {"count": 0}, + "function (obj, prev) { prev.count++; }"] # NOTE maybe we can't count on this ordering being right expected = [{"a": 2, "count": 2}, {"a": None, "count": 1}, {"a": 1, "count": 1}] - self.assertEqual(expected, db.test.group(["a"], {}, - {"count": 0}, - "function (obj, prev) { " - "prev.count++; }", - command=True)) - self.assertEqual(expected, db.test.group(["a"], {}, - {"count": 0}, - "function (obj, prev) { " - "prev.count++; }")) + group_checker(args, expected) + + # modifying finalize + args = [["a"], {}, + {"count": 0}, + "function (obj, prev) { prev.count++; }", + "function(obj){obj.count++;}"] + expected = [{"a": 2, "count": 3}, + {"a": None, "count": 2}, + {"a": 1, "count": 2}] + group_checker(args, expected) + + # returning finalize + args = [["a"], {}, + {"count": 0}, + "function (obj, prev) { prev.count++; }", + "function(obj){ return obj.count;}"] + expected = [2, # a:2 + 1, # a:None + 1] # a:1 + group_checker(args, expected) self.assertRaises(OperationFailure, db.test.group, [], {}, {}, "5 ++ 5") self.assertRaises(OperationFailure, db.test.group, [], {}, {}, "5 ++ 5", command=True)