Merge branch 'master' of git://github.com/RedBeard0531/mongo-python-d…

…river into redbeard/master

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):

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)