From de72eddf69e5ea6dbd3591ddfef5602b5b455942 Mon Sep 17 00:00:00 2001
From: Raymond Hettinger <python@rcn.com>
Date: Sun, 13 Jun 2004 15:36:56 +0000
Subject: [PATCH] Fixup error exits in nlargest() and nsmallest().

---
 Modules/_heapqmodule.c | 24 ++++++++++++++++--------
 1 file changed, 16 insertions(+), 8 deletions(-)

diff --git a/Modules/_heapqmodule.c b/Modules/_heapqmodule.c
index 61ee4134392..8074bb111d3 100644
--- a/Modules/_heapqmodule.c
+++ b/Modules/_heapqmodule.c
@@ -230,13 +230,17 @@ nlargest(PyObject *self, PyObject *args)
 		return NULL;
 
 	heap = PyList_New(0);
-	if (it == NULL)
+	if (heap == NULL)
 		goto fail;
 
 	for (i=0 ; i<n ; i++ ){
 		elem = PyIter_Next(it);
-		if (elem == NULL)
-			goto sortit;
+		if (elem == NULL) {
+			if (PyErr_Occurred())
+				goto fail;
+			else
+				goto sortit;
+		}
 		if (PyList_Append(heap, elem) == -1) {
 			Py_DECREF(elem);
 			goto fail;
@@ -271,11 +275,11 @@ nlargest(PyObject *self, PyObject *args)
 		sol = PyList_GET_ITEM(heap, 0);
 	}
 sortit:
-	Py_DECREF(it);
 	if (PyList_Sort(heap) == -1)
 		goto fail;
 	if (PyList_Reverse(heap) == -1)
 		goto fail;
+	Py_DECREF(it);
 	return heap;
 
 fail:
@@ -385,13 +389,17 @@ nsmallest(PyObject *self, PyObject *args)
 		return NULL;
 
 	heap = PyList_New(0);
-	if (it == NULL)
+	if (heap == NULL)
 		goto fail;
 
 	for (i=0 ; i<n ; i++ ){
 		elem = PyIter_Next(it);
-		if (elem == NULL)
-			goto sortit;
+		if (elem == NULL) {
+			if (PyErr_Occurred())
+				goto fail;
+			else
+				goto sortit;
+		}
 		if (PyList_Append(heap, elem) == -1) {
 			Py_DECREF(elem);
 			goto fail;
@@ -429,9 +437,9 @@ nsmallest(PyObject *self, PyObject *args)
 	}
 
 sortit:
-	Py_DECREF(it);
 	if (PyList_Sort(heap) == -1)
 		goto fail;
+	Py_DECREF(it);
 	return heap;
 
 fail: