diff options
| -rw-r--r-- | Lib/test/test_itertools.py | 8 | ||||
| -rw-r--r-- | Modules/itertoolsmodule.c | 7 |
2 files changed, 14 insertions, 1 deletions
diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py index aa80910..0f74c62 100644 --- a/Lib/test/test_itertools.py +++ b/Lib/test/test_itertools.py @@ -1,6 +1,7 @@ import unittest from test import test_support from itertools import * +from weakref import proxy import sys import operator import random @@ -382,6 +383,13 @@ class TestBasicOps(unittest.TestCase): t3 = tnew(t1) self.assert_(list(t1) == list(t2) == list(t3) == list('abc')) + # test that tee objects are weak referencable + a, b = tee(xrange(10)) + p = proxy(a) + self.assertEqual(getattr(p, '__class__'), type(b)) + del a + self.assertRaises(ReferenceError, getattr, p, '__class__') + def test_StopIteration(self): self.assertRaises(StopIteration, izip().next) diff --git a/Modules/itertoolsmodule.c b/Modules/itertoolsmodule.c index 4069ea2..677dd49 100644 --- a/Modules/itertoolsmodule.c +++ b/Modules/itertoolsmodule.c @@ -329,6 +329,7 @@ typedef struct { PyObject_HEAD teedataobject *dataobj; int index; + PyObject *weakreflist; } teeobject; static PyTypeObject teedataobject_type; @@ -452,6 +453,7 @@ tee_copy(teeobject *to) Py_INCREF(to->dataobj); newto->dataobj = to->dataobj; newto->index = to->index; + newto->weakreflist = NULL; return (PyObject *)newto; } @@ -476,6 +478,7 @@ tee_fromiterable(PyObject *iterable) goto done; to->dataobj = (teedataobject *)teedataobject_new(it); to->index = 0; + to->weakreflist = NULL; done: Py_XDECREF(it); return (PyObject *)to; @@ -494,6 +497,8 @@ tee_new(PyTypeObject *type, PyObject *args, PyObject *kw) static void tee_dealloc(teeobject *to) { + if (to->weakreflist != NULL) + PyObject_ClearWeakRefs((PyObject *) to); Py_XDECREF(to->dataobj); PyObject_Del(to); } @@ -533,7 +538,7 @@ static PyTypeObject tee_type = { 0, /* tp_traverse */ 0, /* tp_clear */ 0, /* tp_richcompare */ - 0, /* tp_weaklistoffset */ + offsetof(teeobject, weakreflist), /* tp_weaklistoffset */ PyObject_SelfIter, /* tp_iter */ (iternextfunc)tee_next, /* tp_iternext */ tee_methods, /* tp_methods */ |