Fix and test weak referencing of itertools.tee objects.
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 @@
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 @@
PyObject_HEAD
teedataobject *dataobj;
int index;
+ PyObject *weakreflist;
} teeobject;
static PyTypeObject teedataobject_type;
@@ -452,6 +453,7 @@
Py_INCREF(to->dataobj);
newto->dataobj = to->dataobj;
newto->index = to->index;
+ newto->weakreflist = NULL;
return (PyObject *)newto;
}
@@ -476,6 +478,7 @@
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 @@
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 @@
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 */