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);