Fix and test weak referencing of itertools.tee objects.

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)
 

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 */