From: Raymond Hettinger <python@rcn.com>
Date: Sun, 13 Jun 2004 15:36:56 +0000 (+0000)
Subject: Fixup error exits in nlargest() and nsmallest().
X-Git-Tag: v2.4a1~192
X-Git-Url: https://granicus.if.org/sourcecode?a=commitdiff_plain;h=de72eddf69e5ea6dbd3591ddfef5602b5b455942;p=python

Fixup error exits in nlargest() and nsmallest().
---

diff --git a/Modules/_heapqmodule.c b/Modules/_heapqmodule.c
index 61ee413439..8074bb111d 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: