[Python-checkins] r58304 - in python/trunk: Misc/NEWS Objects/enumobject.c

[raymond.hettinger]

Author: raymond.hettinger
Date: Wed Oct  3 23:18:11 2007
New Revision: 58304

Modified:
   python/trunk/Misc/NEWS
   python/trunk/Objects/enumobject.c
Log:
enumerate() is no longer bounded to using sequences shorter than LONG_MAX.  The possibility of overflow was sending some newsgroup posters into a tizzy.

Modified: python/trunk/Misc/NEWS
==============================================================================
--- python/trunk/Misc/NEWS	(original)
+++ python/trunk/Misc/NEWS	Wed Oct  3 23:18:11 2007
@@ -12,6 +12,10 @@
 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

Modified: python/trunk/Objects/enumobject.c
==============================================================================
--- python/trunk/Objects/enumobject.c	(original)
+++ python/trunk/Objects/enumobject.c	Wed Oct  3 23:18:11 2007
@@ -7,6 +7,7 @@
 	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 @@
 		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 @@
 	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,10 +54,53 @@
 {
 	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)
 {
 	PyObject *next_index;
@@ -62,16 +108,13 @@
 	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);