Fix and test weak referencing of itertools.tee objects.

diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py
index aa809105479489dd87be6bb34c9928263a3ca92a..0f74c6207359c968699438f57c907dbf64008c9b 100644 (file)
--- 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 4069ea2dafbeaeaafca89edb572ff1458f2ff06c..677dd498407c74bf7473544c39fa4ca5218d829b 100644 (file)
--- 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 */