Skip to content

Commit b65fb33

Browse files
committed
SF patch 1546297 (with some tweaks):
Create a real zip iterator object; not using itertools.izip (Brian Holmes).
1 parent e2e23ef commit b65fb33

File tree

3 files changed

+177
-15
lines changed

3 files changed

+177
-15
lines changed

Include/iterobject.h

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,9 @@ PyAPI_DATA(PyTypeObject) PyCallIter_Type;
1616
#define PyCallIter_Check(op) ((op)->ob_type == &PyCallIter_Type)
1717

1818
PyAPI_FUNC(PyObject *) PyCallIter_New(PyObject *, PyObject *);
19+
20+
PyObject* _PyZip_CreateIter(PyObject* args);
21+
1922
#ifdef __cplusplus
2023
}
2124
#endif

Objects/iterobject.c

Lines changed: 171 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -230,3 +230,174 @@ PyTypeObject PyCallIter_Type = {
230230
(iternextfunc)calliter_iternext, /* tp_iternext */
231231
0, /* tp_methods */
232232
};
233+
234+
235+
/*********************** Zip Iterator **************************/
236+
/* Largely copied from itertools.c by Brian Holmes */
237+
238+
typedef struct zipiterobject_t {
239+
PyObject_HEAD
240+
PyTupleObject *it_tuple; /* Set to NULL when iterator is exhausted */
241+
Py_ssize_t resultsize;
242+
PyTupleObject *result; /* Reusable tuple for optimization */
243+
} zipiterobject;
244+
245+
static PyTypeObject PyZipIter_Type; /* Forward */
246+
247+
PyObject *
248+
_PyZip_CreateIter(PyObject* args)
249+
{
250+
Py_ssize_t i;
251+
Py_ssize_t tuplesize;
252+
PyObject* ziptuple;
253+
PyObject* result;
254+
struct zipiterobject_t* zipiter;
255+
256+
assert(PyTuple_Check(args));
257+
258+
if (PyZipIter_Type.ob_type == NULL) {
259+
if (PyType_Ready(&PyZipIter_Type) < 0)
260+
return NULL;
261+
}
262+
263+
tuplesize = PySequence_Length((PyObject*) args);
264+
265+
ziptuple = PyTuple_New(tuplesize);
266+
if (ziptuple == NULL)
267+
return NULL;
268+
269+
for (i = 0; i < tuplesize; i++) {
270+
PyObject *o = PyTuple_GET_ITEM(args, i);
271+
PyObject *it = PyObject_GetIter(o);
272+
if (it == NULL) {
273+
/* XXX Should we do this?
274+
if (PyErr_ExceptionMatches(PyExc_TypeError))
275+
PyErr_Format(PyExc_TypeError,
276+
"zip argument #%zd must support iteration",
277+
I+1);
278+
*/
279+
Py_DECREF(ziptuple);
280+
return NULL;
281+
}
282+
PyTuple_SET_ITEM(ziptuple, i, it);
283+
}
284+
285+
/* create a reusable result holder */
286+
result = PyTuple_New(tuplesize);
287+
if (result == NULL) {
288+
Py_DECREF(ziptuple);
289+
return NULL;
290+
}
291+
for (i = 0; i < tuplesize; i++) {
292+
Py_INCREF(Py_None);
293+
PyTuple_SET_ITEM(result, i, Py_None);
294+
}
295+
296+
zipiter = PyObject_GC_New(zipiterobject, &PyZipIter_Type);
297+
if (zipiter == NULL) {
298+
Py_DECREF(ziptuple);
299+
Py_DECREF(result);
300+
return NULL;
301+
}
302+
303+
zipiter->result = (PyTupleObject*) result;
304+
zipiter->resultsize = tuplesize;
305+
Py_INCREF(ziptuple);
306+
zipiter->it_tuple = (PyTupleObject *) ziptuple;
307+
_PyObject_GC_TRACK(zipiter);
308+
return (PyObject *)zipiter;
309+
}
310+
311+
static void
312+
zipiter_dealloc(zipiterobject *it)
313+
{
314+
_PyObject_GC_UNTRACK(it);
315+
Py_XDECREF(it->it_tuple);
316+
Py_XDECREF(it->result);
317+
PyObject_GC_Del(it);
318+
}
319+
320+
static int
321+
zipiter_traverse(zipiterobject *it, visitproc visit, void *arg)
322+
{
323+
Py_VISIT(it->it_tuple);
324+
Py_VISIT(it->result);
325+
return 0;
326+
}
327+
328+
static PyObject *
329+
zipiter_next(zipiterobject *zit)
330+
{
331+
Py_ssize_t i;
332+
Py_ssize_t tuplesize = zit->resultsize;
333+
PyObject *result = (PyObject*) zit->result;
334+
PyObject *olditem;
335+
336+
if (tuplesize == 0)
337+
return NULL;
338+
339+
if (result->ob_refcnt == 1) {
340+
Py_INCREF(result);
341+
for (i = 0; i < tuplesize; i++) {
342+
PyObject *it = PyTuple_GET_ITEM(zit->it_tuple, i);
343+
assert(PyIter_Check(it));
344+
PyObject *item = (*it->ob_type->tp_iternext)(it);
345+
if (item == NULL) {
346+
Py_DECREF(result);
347+
return NULL;
348+
}
349+
olditem = PyTuple_GET_ITEM(result, i);
350+
PyTuple_SET_ITEM(result, i, item);
351+
Py_DECREF(olditem);
352+
}
353+
} else {
354+
result = PyTuple_New(tuplesize);
355+
if (result == NULL)
356+
return NULL;
357+
for (i = 0; i < tuplesize; i++) {
358+
PyObject *it = PyTuple_GET_ITEM(zit->it_tuple, i);
359+
assert(PyIter_Check(it));
360+
PyObject *item = (*it->ob_type->tp_iternext)(it);
361+
if (item == NULL) {
362+
Py_DECREF(result);
363+
return NULL;
364+
}
365+
PyTuple_SET_ITEM(result, i, item);
366+
}
367+
}
368+
return result;
369+
}
370+
371+
static PyTypeObject PyZipIter_Type = {
372+
PyObject_HEAD_INIT(0)
373+
0, /* ob_size */
374+
"zipiterator", /* tp_name */
375+
sizeof(zipiterobject), /* tp_basicsize */
376+
0, /* tp_itemsize */
377+
/* methods */
378+
(destructor)zipiter_dealloc, /* tp_dealloc */
379+
0, /* tp_print */
380+
0, /* tp_getattr */
381+
0, /* tp_setattr */
382+
0, /* tp_compare */
383+
0, /* tp_repr */
384+
0, /* tp_as_number */
385+
0, /* tp_as_sequence */
386+
0, /* tp_as_mapping */
387+
0, /* tp_hash */
388+
0, /* tp_call */
389+
0, /* tp_str */
390+
PyObject_GenericGetAttr, /* tp_getattro */
391+
0, /* tp_setattro */
392+
0, /* tp_as_buffer */
393+
Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,/* tp_flags */
394+
0, /* tp_doc */
395+
(traverseproc)zipiter_traverse, /* tp_traverse */
396+
0, /* tp_clear */
397+
0, /* tp_richcompare */
398+
0, /* tp_weakzipoffset */
399+
PyObject_SelfIter, /* tp_iter */
400+
(iternextfunc)zipiter_next, /* tp_iternext */
401+
0, /* tp_methods */
402+
0, /* tp_members */
403+
};

Python/bltinmodule.c

Lines changed: 3 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1855,22 +1855,10 @@ is a shortcut for issubclass(X, A) or issubclass(X, B) or ... (etc.).");
18551855
static PyObject*
18561856
builtin_zip(PyObject *self, PyObject *args)
18571857
{
1858-
PyObject *itertools = NULL, *izip = NULL, *result = NULL;
1858+
/* args must be a tuple */
1859+
assert(PyTuple_Check(args));
18591860

1860-
itertools = PyImport_ImportModule("itertools");
1861-
if (itertools == NULL)
1862-
return NULL;
1863-
1864-
izip = PyObject_GetAttrString(itertools, "izip");
1865-
if (izip == NULL)
1866-
goto done;
1867-
1868-
result = PyObject_Call(izip, args, NULL);
1869-
1870-
done:
1871-
Py_XDECREF(itertools);
1872-
Py_XDECREF(izip);
1873-
return result;
1861+
return _PyZip_CreateIter(args);
18741862
}
18751863

18761864

0 commit comments

Comments
 (0)