From 8f6693701c1899168f5bca96549ae2f2f590624e Mon Sep 17 00:00:00 2001
From: Raymond Hettinger <python@rcn.com>
Date: Wed, 3 Oct 2007 21:18:11 +0000
Subject: [PATCH] enumerate() is no longer bounded to using sequences shorter
 than LONG_MAX.  The possibility of overflow was sending some newsgroup
 posters into a tizzy.

---
 Misc/NEWS            |  4 ++++
 Objects/enumobject.c | 55 +++++++++++++++++++++++++++++++++++++++-----
 2 files changed, 53 insertions(+), 6 deletions(-)

diff --git a/Misc/NEWS b/Misc/NEWS
index b01633abda3..01f860bc3ff 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -12,6 +12,10 @@ What's New in Python 2.6 alpha 1?
 Core and builtins
 -----------------
 
+- The enumerate() builtin function is no longer bounded to sequences smaller
+  than LONG_MAX.  Formerly, it raised an OverflowError.  Now, automatically
+  shifts from ints to longs.
+
 - Issue #1686386: Tuple's tp_repr did not take into account the possibility of
   having a self-referential tuple, which is possible from C code.  Nor did
   object's tp_str consider that a type's tp_str could do something that could
diff --git a/Objects/enumobject.c b/Objects/enumobject.c
index 6997fdccf70..0a3694d0024 100644
--- a/Objects/enumobject.c
+++ b/Objects/enumobject.c
@@ -7,6 +7,7 @@ typedef struct {
 	long      en_index;        /* current index of enumeration */
 	PyObject* en_sit;          /* secondary iterator of enumeration */
 	PyObject* en_result;	   /* result tuple  */
+	PyObject* en_longindex;	   /* index for sequences >= LONG_MAX */
 } enumobject;
 
 static PyObject *
@@ -25,6 +26,7 @@ enum_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
 		return NULL;
 	en->en_index = 0;
 	en->en_sit = PyObject_GetIter(seq);
+	en->en_longindex = NULL;
 	if (en->en_sit == NULL) {
 		Py_DECREF(en);
 		return NULL;
@@ -43,6 +45,7 @@ enum_dealloc(enumobject *en)
 	PyObject_GC_UnTrack(en);
 	Py_XDECREF(en->en_sit);
 	Py_XDECREF(en->en_result);
+	Py_XDECREF(en->en_longindex);
 	Py_Type(en)->tp_free(en);
 }
 
@@ -51,9 +54,52 @@ enum_traverse(enumobject *en, visitproc visit, void *arg)
 {
 	Py_VISIT(en->en_sit);
 	Py_VISIT(en->en_result);
+	Py_VISIT(en->en_longindex);
 	return 0;
 }
 
+static PyObject *
+enum_next_long(enumobject *en, PyObject* next_item)
+{
+	static PyObject *one = NULL;
+	PyObject *result = en->en_result;
+	PyObject *next_index;
+	PyObject *stepped_up;
+
+	if (en->en_longindex == NULL) {
+		en->en_longindex = PyInt_FromLong(LONG_MAX);
+		if (en->en_longindex == NULL)
+			return NULL;
+	}
+	if (one == NULL) {
+		one = PyInt_FromLong(1);
+		if (one == NULL)
+			return NULL;
+	}
+	next_index = en->en_longindex;
+	assert(next_index != NULL);
+	stepped_up = PyNumber_Add(next_index, one);
+	if (stepped_up == NULL)
+		return NULL;
+	en->en_longindex = stepped_up;
+
+	if (result->ob_refcnt == 1) {
+		Py_INCREF(result);
+		Py_DECREF(PyTuple_GET_ITEM(result, 0));
+		Py_DECREF(PyTuple_GET_ITEM(result, 1));
+	} else {
+		result = PyTuple_New(2);
+		if (result == NULL) {
+			Py_DECREF(next_index);
+			Py_DECREF(next_item);
+			return NULL;
+		}
+	}
+	PyTuple_SET_ITEM(result, 0, next_index);
+	PyTuple_SET_ITEM(result, 1, next_item);
+	return result;
+}
+
 static PyObject *
 enum_next(enumobject *en)
 {
@@ -62,16 +108,13 @@ enum_next(enumobject *en)
 	PyObject *result = en->en_result;
 	PyObject *it = en->en_sit;
 
-	if (en->en_index == LONG_MAX) {
-		PyErr_SetString(PyExc_OverflowError,
-			"enumerate() is limited to LONG_MAX items");                
-		return NULL;         
-	}
-
 	next_item = (*Py_Type(it)->tp_iternext)(it);
 	if (next_item == NULL)
 		return NULL;
 
+	if (en->en_index == LONG_MAX)
+		return enum_next_long(en, next_item);
+
 	next_index = PyInt_FromLong(en->en_index);
 	if (next_index == NULL) {
 		Py_DECREF(next_item);