From 70ad67bd5efa74f1daea9f3ea259f76af381da40 Mon Sep 17 00:00:00 2001 From: Mike Dirolf Date: Fri, 18 Dec 2009 13:09:12 -0500 Subject: [PATCH] support keyf if key is a string and not a list. all groups are run as commands now --- pymongo/collection.py | 103 +++++++++++++--------------------------- test/test_collection.py | 75 ++++++++++++++--------------- 2 files changed, 68 insertions(+), 110 deletions(-) diff --git a/pymongo/collection.py b/pymongo/collection.py index 51956a3e4..4d874eff6 100644 --- a/pymongo/collection.py +++ b/pymongo/collection.py @@ -625,16 +625,26 @@ class Collection(object): return options - def group(self, keys, condition, initial, reduce, finalize=None, + # TODO key and condition ought to be optional, but deprecation + # could be painful as argument order would have to change. + def group(self, key, condition, initial, reduce, finalize=None, command=True): - """Perform a query similar to an SQL group by operation. + """Perform a query similar to an SQL *group by* operation. Returns an array of grouped items. + The `key` parameter can be: + + - ``None`` to use the entire document as a key. + - A :class:`list` of keys (each a :class:`basestring`) to group by. + - A :class:`basestring` or :class:`~pymongo.code.Code` instance + containing a JavaScript function to be applied to each document, + returning the key to group by. + :Parameters: - - `keys`: list of fields to group by - - `condition`: specification of rows to be considered (as a `find` - query specification) + - `key`: fields to group by (see above description) + - `condition`: specification of rows to be + considered (as a :meth:`find` 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. @@ -644,75 +654,28 @@ class Collection(object): .. versionchanged:: 1.3 The `command` argument now defaults to ``True`` and is deprecated. + .. versionchanged:: 1.3+ + The `key` argument can now be ``None`` or a JavaScript function, + in addition to a :class:`list` of keys. """ - - #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.", + if not command: + warnings.warn("eval-based groups are deprecated, and the " + "command option will be removed.", DeprecationWarning) - if command: - if not isinstance(reduce, Code): - reduce = Code(reduce) - group = {"ns": self.__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"] + group = {} + if isinstance(key, basestring): + group["$keyf"] = Code(key) + elif key is not None: + group = {"key": self._fields_list_to_dict(key)} + group["ns"] = self.__name + group["$reduce"] = Code(reduce) + group["cond"] = condition + group["initial"] = initial + if finalize is not None: + group["finalize"] = Code(finalize) - warnings.warn("eval-based groups are deprecated. please use command=True.", - DeprecationWarning) - - scope = {} - if isinstance(reduce, Code): - scope = reduce.scope - scope.update({"ns": self.__name, - "keys": keys, - "condition": condition, - "initial": initial}) - - group_function = """function () { - 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(); - - var key = {}; - for (var i = 0; i < keys.length; i++) { - var k = keys[i]; - key[k] = obj[k]; - } - - var aggObj = map.get(key); - if (aggObj == null) { - var newObj = Object.extend({}, key); - aggObj = Object.extend(newObj, initial); - map.put(key, aggObj); - } - reduce_function(obj, aggObj); - } - - 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"] + return self.__database._command({"group":group})["retval"] def rename(self, new_name): """Rename this collection. diff --git a/test/test_collection.py b/test/test_collection.py index e973c4664..ddeaea83b 100644 --- a/test/test_collection.py +++ b/test/test_collection.py @@ -594,59 +594,54 @@ class TestCollection(unittest.TestCase): self.assertEqual(eval, expected) - args = [[], {}, - {"count": 0}, - "function (obj, prev) { prev.count++; }"] - expected = [] - group_checker(args, expected) + self.assertEqual([], db.test.group([], {}, {"count": 0}, + "function (obj, prev) { prev.count++; }")) db.test.save({"a": 2}) db.test.save({"b": 5}) db.test.save({"a": 1}) - args = [[], {}, - {"count": 0}, - "function (obj, prev) { prev.count++; }"] - expected = [{'count': 3}] - group_checker(args, expected) + self.assertEqual([{"count": 3}], + db.test.group([], {}, {"count": 0}, + "function (obj, prev) { prev.count++; }")) - args = [[], - {"a": {"$gt": 1}}, - {"count": 0}, - "function (obj, prev) { prev.count++; }"] - expected = [{'count': 1}] - group_checker(args, expected) + self.assertEqual([{"count": 1}], + db.test.group([], {"a": {"$gt": 1}}, {"count": 0}, + "function (obj, prev) { prev.count++; }")) 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}] - group_checker(args, expected) + self.assertEqual([{"a": 2, "count": 2}, + {"a": None, "count": 1}, + {"a": 1, "count": 1}], + db.test.group(["a"], {}, {"count": 0}, + "function (obj, prev) { prev.count++; }")) # 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) + self.assertEqual([{"a": 2, "count": 3}, + {"a": None, "count": 2}, + {"a": 1, "count": 2}], + db.test.group(["a"], {}, {"count": 0}, + "function (obj, prev) { prev.count++; }", + "function (obj) { obj.count++; }")) # 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.assertEqual([2, 1, 1], + db.test.group(["a"], {}, {"count": 0}, + "function (obj, prev) { prev.count++; }", + "function (obj) { return obj.count; }")) + + # keyf + self.assertEqual([2, 2], + db.test.group("function (obj) { if (obj.a == 2) { return {a: true} }; " + "return {b: true}; }", {}, {"count": 0}, + "function (obj, prev) { prev.count++; }", + "function (obj) { return obj.count; }")) + + # no key + self.assertEqual([{"count": 4}], + db.test.group(None, {}, {"count": 0}, + "function (obj, prev) { prev.count++; }")) warnings.simplefilter("error") self.assertRaises(DeprecationWarning,