Merge branch 'master' of git://github.com/RedBeard0531/mongo-python-driver into redbeard/master

Author: Mike Dirolf

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)