diff options
| -rw-r--r-- | Include/iterobject.h | 3 | ||||
| -rw-r--r-- | Objects/iterobject.c | 171 | ||||
| -rw-r--r-- | Python/bltinmodule.c | 18 |
3 files changed, 177 insertions, 15 deletions
diff --git a/Include/iterobject.h b/Include/iterobject.h index c078ebb..69deb45 100644 --- a/Include/iterobject.h +++ b/Include/iterobject.h @@ -16,6 +16,9 @@ PyAPI_DATA(PyTypeObject) PyCallIter_Type; #define PyCallIter_Check(op) ((op)->ob_type == &PyCallIter_Type) PyAPI_FUNC(PyObject *) PyCallIter_New(PyObject *, PyObject *); + +PyObject* _PyZip_CreateIter(PyObject* args); + #ifdef __cplusplus } #endif diff --git a/Objects/iterobject.c b/Objects/iterobject.c index cf839f4..dc4c583 100644 --- a/Objects/iterobject.c +++ b/Objects/iterobject.c @@ -230,3 +230,174 @@ PyTypeObject PyCallIter_Type = { (iternextfunc)calliter_iternext, /* tp_iternext */ 0, /* tp_methods */ }; + + +/*********************** Zip Iterator **************************/ +/* Largely copied from itertools.c by Brian Holmes */ + +typedef struct zipiterobject_t { + PyObject_HEAD + PyTupleObject *it_tuple; /* Set to NULL when iterator is exhausted */ + Py_ssize_t resultsize; + PyTupleObject *result; /* Reusable tuple for optimization */ +} zipiterobject; + +static PyTypeObject PyZipIter_Type; /* Forward */ + +PyObject * +_PyZip_CreateIter(PyObject* args) +{ + Py_ssize_t i; + Py_ssize_t tuplesize; + PyObject* ziptuple; + PyObject* result; + struct zipiterobject_t* zipiter; + + assert(PyTuple_Check(args)); + + if (PyZipIter_Type.ob_type == NULL) { + if (PyType_Ready(&PyZipIter_Type) < 0) + return NULL; + } + + tuplesize = PySequence_Length((PyObject*) args); + + ziptuple = PyTuple_New(tuplesize); + if (ziptuple == NULL) + return NULL; + + for (i = 0; i < tuplesize; i++) { + PyObject *o = PyTuple_GET_ITEM(args, i); + PyObject *it = PyObject_GetIter(o); + if (it == NULL) { + /* XXX Should we do this? + if (PyErr_ExceptionMatches(PyExc_TypeError)) + PyErr_Format(PyExc_TypeError, + "zip argument #%zd must support iteration", + I+1); + */ + Py_DECREF(ziptuple); + return NULL; + } + PyTuple_SET_ITEM(ziptuple, i, it); + } + + /* create a reusable result holder */ + result = PyTuple_New(tuplesize); + if (result == NULL) { + Py_DECREF(ziptuple); + return NULL; + } + for (i = 0; i < tuplesize; i++) { + Py_INCREF(Py_None); + PyTuple_SET_ITEM(result, i, Py_None); + } + + zipiter = PyObject_GC_New(zipiterobject, &PyZipIter_Type); + if (zipiter == NULL) { + Py_DECREF(ziptuple); + Py_DECREF(result); + return NULL; + } + + zipiter->result = (PyTupleObject*) result; + zipiter->resultsize = tuplesize; + Py_INCREF(ziptuple); + zipiter->it_tuple = (PyTupleObject *) ziptuple; + _PyObject_GC_TRACK(zipiter); + return (PyObject *)zipiter; +} + +static void +zipiter_dealloc(zipiterobject *it) +{ + _PyObject_GC_UNTRACK(it); + Py_XDECREF(it->it_tuple); + Py_XDECREF(it->result); + PyObject_GC_Del(it); +} + +static int +zipiter_traverse(zipiterobject *it, visitproc visit, void *arg) +{ + Py_VISIT(it->it_tuple); + Py_VISIT(it->result); + return 0; +} + +static PyObject * +zipiter_next(zipiterobject *zit) +{ + Py_ssize_t i; + Py_ssize_t tuplesize = zit->resultsize; + PyObject *result = (PyObject*) zit->result; + PyObject *olditem; + + if (tuplesize == 0) + return NULL; + + if (result->ob_refcnt == 1) { + Py_INCREF(result); + for (i = 0; i < tuplesize; i++) { + PyObject *it = PyTuple_GET_ITEM(zit->it_tuple, i); + assert(PyIter_Check(it)); + PyObject *item = (*it->ob_type->tp_iternext)(it); + if (item == NULL) { + Py_DECREF(result); + return NULL; + } + olditem = PyTuple_GET_ITEM(result, i); + PyTuple_SET_ITEM(result, i, item); + Py_DECREF(olditem); + } + } else { + result = PyTuple_New(tuplesize); + if (result == NULL) + return NULL; + for (i = 0; i < tuplesize; i++) { + PyObject *it = PyTuple_GET_ITEM(zit->it_tuple, i); + assert(PyIter_Check(it)); + PyObject *item = (*it->ob_type->tp_iternext)(it); + if (item == NULL) { + Py_DECREF(result); + return NULL; + } + PyTuple_SET_ITEM(result, i, item); + } + } + return result; +} + +static PyTypeObject PyZipIter_Type = { + PyObject_HEAD_INIT(0) + 0, /* ob_size */ + "zipiterator", /* tp_name */ + sizeof(zipiterobject), /* tp_basicsize */ + 0, /* tp_itemsize */ + /* methods */ + (destructor)zipiter_dealloc, /* tp_dealloc */ + 0, /* tp_print */ + 0, /* tp_getattr */ + 0, /* tp_setattr */ + 0, /* tp_compare */ + 0, /* tp_repr */ + 0, /* tp_as_number */ + 0, /* tp_as_sequence */ + 0, /* tp_as_mapping */ + 0, /* tp_hash */ + 0, /* tp_call */ + 0, /* tp_str */ + PyObject_GenericGetAttr, /* tp_getattro */ + 0, /* tp_setattro */ + 0, /* tp_as_buffer */ + Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,/* tp_flags */ + 0, /* tp_doc */ + (traverseproc)zipiter_traverse, /* tp_traverse */ + 0, /* tp_clear */ + 0, /* tp_richcompare */ + 0, /* tp_weakzipoffset */ + PyObject_SelfIter, /* tp_iter */ + (iternextfunc)zipiter_next, /* tp_iternext */ + 0, /* tp_methods */ + 0, /* tp_members */ +}; diff --git a/Python/bltinmodule.c b/Python/bltinmodule.c index 6ca2a28..200ec26 100644 --- a/Python/bltinmodule.c +++ b/Python/bltinmodule.c @@ -1855,22 +1855,10 @@ is a shortcut for issubclass(X, A) or issubclass(X, B) or ... (etc.)."); static PyObject* builtin_zip(PyObject *self, PyObject *args) { - PyObject *itertools = NULL, *izip = NULL, *result = NULL; + /* args must be a tuple */ + assert(PyTuple_Check(args)); - itertools = PyImport_ImportModule("itertools"); - if (itertools == NULL) - return NULL; - - izip = PyObject_GetAttrString(itertools, "izip"); - if (izip == NULL) - goto done; - - result = PyObject_Call(izip, args, NULL); - - done: - Py_XDECREF(itertools); - Py_XDECREF(izip); - return result; + return _PyZip_CreateIter(args); } |