SF patch #1077353: add key= argument to min and max

(First draft of patch contributed by Steven Bethard.)

5 files changed, 164 insertions, 45 deletions

diff --git a/Doc/lib/libfuncs.tex b/Doc/lib/libfuncs.tex
index 2099f3f..904b4e1 100644
--- a/Doc/lib/libfuncs.tex
+++ b/Doc/lib/libfuncs.tex
@@ -642,16 +642,28 @@ class C:
   of sequence; the result is always a list.
 \end{funcdesc}
 
-\begin{funcdesc}{max}{s\optional{, args...}}
+\begin{funcdesc}{max}{s\optional{, args...}\optional{key}}
   With a single argument \var{s}, return the largest item of a
   non-empty sequence (such as a string, tuple or list).  With more
   than one argument, return the largest of the arguments.
+
+  The optional \var{key} argument specifies a one argument ordering
+  function like that used for \method{list.sort()}.  The \var{key}
+  argument, if supplied, must be in keyword form (for example,
+  \samp{max(a,b,c,key=func)}).
+  \versionchanged[Added support for the optional \var{key} argument]{2.5}
 \end{funcdesc}
 
 \begin{funcdesc}{min}{s\optional{, args...}}
   With a single argument \var{s}, return the smallest item of a
   non-empty sequence (such as a string, tuple or list).  With more
   than one argument, return the smallest of the arguments.
+
+  The optional \var{key} argument specifies a one argument ordering
+  function like that used for \method{list.sort()}.  The \var{key}
+  argument, if supplied, must be in keyword form (for example,
+  \samp{min(a,b,c,key=func)}).
+  \versionchanged[Added support for the optional \var{key} argument]{2.5}           
 \end{funcdesc}
 
 \begin{funcdesc}{object}{}
diff --git a/Lib/test/test_builtin.py b/Lib/test/test_builtin.py
index f023b11..1dae603 100644
--- a/Lib/test/test_builtin.py
+++ b/Lib/test/test_builtin.py
@@ -1,7 +1,8 @@
 # Python test set -- built-in functions
 
 import test.test_support, unittest
-from test.test_support import fcmp, have_unicode, TESTFN, unlink
+from test.test_support import fcmp, have_unicode, TESTFN, unlink, run_unittest
+from operator import neg
 
 import sys, warnings, cStringIO, random, UserDict
 warnings.filterwarnings("ignore", "hex../oct.. of negative int",
@@ -9,6 +10,10 @@ warnings.filterwarnings("ignore", "hex../oct.. of negative int",
 warnings.filterwarnings("ignore", "integer argument expected",
                         DeprecationWarning, "unittest")
 
+# count the number of test runs.
+# used to skip running test_execfile() multiple times
+numruns = 0
+
 class Squares:
 
     def __init__(self, max):
@@ -343,6 +348,11 @@ class BuiltinTest(unittest.TestCase):
     execfile(TESTFN)
 
     def test_execfile(self):
+        global numruns
+        if numruns:
+            return
+        numruns += 1
+        
         globals = {'a': 1, 'b': 2}
         locals = {'b': 200, 'c': 300}
 
@@ -845,6 +855,30 @@ class BuiltinTest(unittest.TestCase):
         self.assertEqual(max(1L, 2.0, 3), 3)
         self.assertEqual(max(1.0, 2, 3L), 3L)
 
+        for stmt in (
+            "max(key=int)",                 # no args
+            "max(1, key=int)",              # single arg not iterable
+            "max(1, 2, keystone=int)",      # wrong keyword
+            "max(1, 2, key=int, abc=int)",  # two many keywords
+            "max(1, 2, key=1)",             # keyfunc is not callable
+            ):
+                try:
+                    exec(stmt) in globals()
+                except TypeError:
+                    pass
+                else:
+                    self.fail(stmt)
+
+        self.assertEqual(max((1,), key=neg), 1)     # one elem iterable
+        self.assertEqual(max((1,2), key=neg), 1)    # two elem iterable
+        self.assertEqual(max(1, 2, key=neg), 1)     # two elems
+
+        data = [random.randrange(200) for i in range(100)]
+        keys = dict((elem, random.randrange(50)) for elem in data)
+        f = keys.__getitem__
+        self.assertEqual(max(data, key=f),
+                         sorted(reversed(data), key=f)[-1])
+
     def test_min(self):
         self.assertEqual(min('123123'), '1')
         self.assertEqual(min(1, 2, 3), 1)
@@ -867,6 +901,30 @@ class BuiltinTest(unittest.TestCase):
                 raise ValueError
         self.assertRaises(ValueError, min, (42, BadNumber()))
 
+        for stmt in (
+            "min(key=int)",                 # no args
+            "min(1, key=int)",              # single arg not iterable
+            "min(1, 2, keystone=int)",      # wrong keyword
+            "min(1, 2, key=int, abc=int)",  # two many keywords
+            "min(1, 2, key=1)",             # keyfunc is not callable
+            ):
+                try:
+                    exec(stmt) in globals()
+                except TypeError:
+                    pass
+                else:
+                    self.fail(stmt)
+
+        self.assertEqual(min((1,), key=neg), 1)     # one elem iterable
+        self.assertEqual(min((1,2), key=neg), 2)    # two elem iterable
+        self.assertEqual(min(1, 2, key=neg), 2)     # two elems
+
+        data = [random.randrange(200) for i in range(100)]
+        keys = dict((elem, random.randrange(50)) for elem in data)
+        f = keys.__getitem__
+        self.assertEqual(min(data, key=f),
+                         sorted(data, key=f)[0])
+
     def test_oct(self):
         self.assertEqual(oct(100), '0144')
         self.assertEqual(oct(100L), '0144L')
@@ -1313,8 +1371,21 @@ class TestSorted(unittest.TestCase):
         data = 'The quick Brown fox Jumped over The lazy Dog'.split()
         self.assertRaises(TypeError, sorted, data, None, lambda x,y: 0)
 
-def test_main():
-    test.test_support.run_unittest(BuiltinTest, TestSorted)
+def test_main(verbose=None):
+    test_classes = (BuiltinTest, TestSorted)
+
+    run_unittest(*test_classes)
+
+    # verify reference counting
+    if verbose and hasattr(sys, "gettotalrefcount"):
+        import gc
+        counts = [None] * 5
+        for i in xrange(len(counts)):
+            run_unittest(*test_classes)
+            gc.collect()
+            counts[i] = sys.gettotalrefcount()
+        print counts
+
 
 if __name__ == "__main__":
-    test_main()
+    test_main(verbose=True)
diff --git a/Misc/ACKS b/Misc/ACKS
index 2c8038f..c83d479 100644
--- a/Misc/ACKS
+++ b/Misc/ACKS
@@ -52,6 +52,7 @@ Alexander Belopolsky
 Andy Bensky
 Michel Van den Bergh
 Eric Beser
+Steven Bethard
 Stephen Bevan
 Ron Bickers
 Dominic Binks
diff --git a/Misc/NEWS b/Misc/NEWS
index 4029c1d..a1dfa45 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -10,6 +10,9 @@ What's New in Python 2.5 alpha 1?
 Core and builtins
 -----------------
 
+- min() and max() now support key= arguments with the same meaning as in
+  list.sort().
+
 
 Extension Modules
 -----------------
@@ -19,7 +22,7 @@ Library
 -------
 
 - heapq.nsmallest() and heapq.nlargest() now support key= arguments with
-  the same meaning as for list.sort().
+  the same meaning as in list.sort().
 
 
 Build
diff --git a/Python/bltinmodule.c b/Python/bltinmodule.c
index 6fbe799..e189952 100644
--- a/Python/bltinmodule.c
+++ b/Python/bltinmodule.c
@@ -1114,82 +1114,114 @@ Update and return a dictionary containing the current scope's local variables.")
 
 
 static PyObject *
-min_max(PyObject *args, int op)
+min_max(PyObject *args, PyObject *kwds, int op)
 {
+	PyObject *v, *it, *item, *val, *maxitem, *maxval, *keyfunc=NULL;
 	const char *name = op == Py_LT ? "min" : "max";
-	PyObject *v, *w, *x, *it;
 
 	if (PyTuple_Size(args) > 1)
 		v = args;
 	else if (!PyArg_UnpackTuple(args, (char *)name, 1, 1, &v))
 		return NULL;
 
+	if (kwds != NULL && PyDict_Check(kwds) && PyDict_Size(kwds)) {
+		keyfunc = PyDict_GetItemString(kwds, "key");
+		if (PyDict_Size(kwds)!=1  ||  keyfunc == NULL) {
+			PyErr_Format(PyExc_TypeError, 
+				"%s() got an unexpected keyword argument", name);
+			return NULL;
+		}
+	} 
+
 	it = PyObject_GetIter(v);
 	if (it == NULL)
 		return NULL;
 
-	w = NULL;  /* the result */
-	for (;;) {
-		x = PyIter_Next(it);
-		if (x == NULL) {
-			if (PyErr_Occurred()) {
-				Py_XDECREF(w);
-				Py_DECREF(it);
-				return NULL;
-			}
-			break;
+	maxitem = NULL; /* the result */
+	maxval = NULL;  /* the value associated with the result */
+	while (item = PyIter_Next(it)) {
+		/* get the value from the key function */
+		if (keyfunc != NULL) {
+			val = PyObject_CallFunctionObjArgs(keyfunc, item, NULL);
+			if (val == NULL)
+				goto Fail_it_item;
+		}
+		/* no key function; the value is the item */
+		else {
+			val = item;
+			Py_INCREF(val);
 		}
 
-		if (w == NULL)
-			w = x;
+		/* maximum value and item are unset; set them */
+		if (maxval == NULL) {
+			maxitem = item;
+			maxval = val;
+		}
+		/* maximum value and item are set; update them as necessary */
 		else {
-			int cmp = PyObject_RichCompareBool(x, w, op);
-			if (cmp > 0) {
-				Py_DECREF(w);
-				w = x;
+			int cmp = PyObject_RichCompareBool(val, maxval, op);
+			if (cmp < 0)
+				goto Fail_it_item_and_val;
+			else if (cmp > 0) {
+				Py_DECREF(maxval);
+				Py_DECREF(maxitem);
+				maxval = val;
+				maxitem = item;
 			}
-			else if (cmp < 0) {
-				Py_DECREF(x);
-				Py_DECREF(w);
-				Py_DECREF(it);
-				return NULL;
+			else {
+				Py_DECREF(item);
+				Py_DECREF(val);
 			}
-			else
-				Py_DECREF(x);
 		}
 	}
-	if (w == NULL)
+	if (PyErr_Occurred())
+		goto Fail_it;
+	if (maxval == NULL) {
 		PyErr_Format(PyExc_ValueError,
 			     "%s() arg is an empty sequence", name);
+		assert(maxitem == NULL);
+	}
+	else
+		Py_DECREF(maxval);
 	Py_DECREF(it);
-	return w;
+	return maxitem;
+
+Fail_it_item_and_val:
+	Py_DECREF(val);
+Fail_it_item:
+	Py_DECREF(item);
+Fail_it:
+	Py_XDECREF(maxval);
+	Py_XDECREF(maxitem);
+	Py_DECREF(it);
+	return NULL;
 }
 
 static PyObject *
-builtin_min(PyObject *self, PyObject *v)
+builtin_min(PyObject *self, PyObject *args, PyObject *kwds)
 {
-	return min_max(v, Py_LT);
+	return min_max(args, kwds, Py_LT);
 }
 
 PyDoc_STRVAR(min_doc,
-"min(sequence) -> value\n\
-min(a, b, c, ...) -> value\n\
+"min(iterable[, key=func]) -> value\n\
+min(a, b, c, ...[, key=func]) -> value\n\
 \n\
-With a single sequence argument, return its smallest item.\n\
+With a single iterable argument, return its smallest item.\n\
 With two or more arguments, return the smallest argument.");
 
 
 static PyObject *
-builtin_max(PyObject *self, PyObject *v)
+builtin_max(PyObject *self, PyObject *args, PyObject *kwds)
 {
-	return min_max(v, Py_GT);
+	return min_max(args, kwds, Py_GT);
 }
 
 PyDoc_STRVAR(max_doc,
-"max(sequence) -> value\n\
-max(a, b, c, ...) -> value\n\
+"max(iterable[, key=func]) -> value\n\
+max(a, b, c, ...[, key=func]) -> value\n\
 \n\
-With a single sequence argument, return its largest item.\n\
+With a single iterable argument, return its largest item.\n\
 With two or more arguments, return the largest argument.");
 
 
@@ -2119,8 +2151,8 @@ static PyMethodDef builtin_methods[] = {
  	{"len",		builtin_len,        METH_O, len_doc},
  	{"locals",	(PyCFunction)builtin_locals,     METH_NOARGS, locals_doc},
  	{"map",		builtin_map,        METH_VARARGS, map_doc},
- 	{"max",		builtin_max,        METH_VARARGS, max_doc},
- 	{"min",		builtin_min,        METH_VARARGS, min_doc},
+ 	{"max",		(PyCFunction)builtin_max,        METH_VARARGS | METH_KEYWORDS, max_doc},
+ 	{"min",		(PyCFunction)builtin_min,        METH_VARARGS | METH_KEYWORDS, min_doc},
  	{"oct",		builtin_oct,        METH_O, oct_doc},
  	{"ord",		builtin_ord,        METH_O, ord_doc},
  	{"pow",		builtin_pow,        METH_VARARGS, pow_doc},