* Beefed-up tests

* Allow tuple re-use
* Call tp_iternext directly

diff --git a/Lib/test/test_enumerate.py b/Lib/test/test_enumerate.py
index 701b0c0b1a214e061a0ac57c68cd5128c93b528a..00f97ef6ffd477a88e64da2f5483a17659f4dd1e 100644 (file)
--- a/Lib/test/test_enumerate.py
+++ b/Lib/test/test_enumerate.py
@@ -1,9 +1,8 @@
 import unittest
+from sets import Set
 
 from test import test_support
 
-seq, res = 'abc', [(0,'a'), (1,'b'), (2,'c')]
-
 class G:
     'Sequence using __getitem__'
     def __init__(self, seqn):
@@ -65,37 +64,49 @@ class N:
 class EnumerateTestCase(unittest.TestCase):
 
     enum = enumerate
+    seq, res = 'abc', [(0,'a'), (1,'b'), (2,'c')]
 
     def test_basicfunction(self):
-        self.assertEqual(type(self.enum(seq)), self.enum)
-        e = self.enum(seq)
+        self.assertEqual(type(self.enum(self.seq)), self.enum)
+        e = self.enum(self.seq)
         self.assertEqual(iter(e), e)
-        self.assertEqual(list(self.enum(seq)), res)
+        self.assertEqual(list(self.enum(self.seq)), self.res)
         self.enum.__doc__
 
     def test_getitemseqn(self):
-        self.assertEqual(list(self.enum(G(seq))), res)
+        self.assertEqual(list(self.enum(G(self.seq))), self.res)
         e = self.enum(G(''))
         self.assertRaises(StopIteration, e.next)
 
     def test_iteratorseqn(self):
-        self.assertEqual(list(self.enum(I(seq))), res)
+        self.assertEqual(list(self.enum(I(self.seq))), self.res)
         e = self.enum(I(''))
         self.assertRaises(StopIteration, e.next)
 
     def test_iteratorgenerator(self):
-        self.assertEqual(list(self.enum(Ig(seq))), res)
+        self.assertEqual(list(self.enum(Ig(self.seq))), self.res)
         e = self.enum(Ig(''))
         self.assertRaises(StopIteration, e.next)
 
     def test_noniterable(self):
-        self.assertRaises(TypeError, self.enum, X(seq))
+        self.assertRaises(TypeError, self.enum, X(self.seq))
 
     def test_illformediterable(self):
-        self.assertRaises(TypeError, list, self.enum(N(seq)))
+        self.assertRaises(TypeError, list, self.enum(N(self.seq)))
 
     def test_exception_propagation(self):
-        self.assertRaises(ZeroDivisionError, list, self.enum(E(seq)))
+        self.assertRaises(ZeroDivisionError, list, self.enum(E(self.seq)))
+
+    def test_argumentcheck(self):
+        self.assertRaises(TypeError, self.enum) # no arguments
+        self.assertRaises(TypeError, self.enum, 1) # wrong type (not iterable)
+        self.assertRaises(TypeError, self.enum, 'abc', 2) # too many arguments
+
+    def test_tuple_reuse(self):
+        # Tests an implementation detail where tuple is reused
+        # whenever nothing else holds a reference to it
+        self.assertEqual(len(Set(map(id, list(self.seq)))), len(self.seq))
+        self.assertEqual(len(Set(map(id, enumerate(self.seq)))), min(1,len(self.seq)))
 
 class MyEnum(enumerate):
     pass
@@ -104,8 +115,28 @@ class SubclassTestCase(EnumerateTestCase):
 
     enum = MyEnum
 
-def test_main():
-    test_support.run_unittest(EnumerateTestCase, SubclassTestCase)
+class TestEmpty(EnumerateTestCase):
+
+    seq, res = '', []
+
+class TestBig(EnumerateTestCase):
+
+    seq = range(10,20000,2)
+    res = zip(range(20000), seq)
+
+
+def test_main(verbose=None):
+    testclasses = (EnumerateTestCase, SubclassTestCase, TestEmpty, TestBig)
+    test_support.run_unittest(*testclasses)
+
+    # verify reference counting
+    import sys
+    if verbose and hasattr(sys, "gettotalrefcount"):
+        counts = [None] * 5
+        for i in xrange(len(counts)):
+            test_support.run_unittest(*testclasses)
+            counts[i] = sys.gettotalrefcount()
+        print counts
 
 if __name__ == "__main__":
-    test_main()
+    test_main(verbose=True)

diff --git a/Objects/enumobject.c b/Objects/enumobject.c
index 66ee870d469fae980989f87eee34e46e8b988470..5fbecbfa9ab6bc6e93862b56f3374eba8c671fe4 100644 (file)
--- a/Objects/enumobject.c
+++ b/Objects/enumobject.c
@@ -6,6 +6,7 @@ typedef struct {
        PyObject_HEAD
        long      en_index;        /* current index of enumeration */
        PyObject* en_sit;          /* secondary iterator of enumeration */
+       PyObject* en_result;       /* result tuple  */
 } enumobject;
 
 PyTypeObject PyEnum_Type;
@@ -30,6 +31,16 @@ enum_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
                Py_DECREF(en);
                return NULL;
        }
+       en->en_result = PyTuple_New(2);
+       if (en->en_result == NULL) {
+               Py_DECREF(en->en_sit);
+               Py_DECREF(en);
+               return NULL;
+       }
+       Py_INCREF(Py_None);
+       PyTuple_SET_ITEM(en->en_result, 0, Py_None);
+       Py_INCREF(Py_None);
+       PyTuple_SET_ITEM(en->en_result, 1, Py_None);
        return (PyObject *)en;
 }
 
@@ -38,42 +49,60 @@ enum_dealloc(enumobject *en)
 {
        PyObject_GC_UnTrack(en);
        Py_XDECREF(en->en_sit);
+       Py_XDECREF(en->en_result);
        en->ob_type->tp_free(en);
 }
 
 static int
 enum_traverse(enumobject *en, visitproc visit, void *arg)
 {
-       if (en->en_sit)
-               return visit(en->en_sit, arg);
+       int err;
+
+       if (en->en_sit) {
+               err = visit(en->en_sit, arg);
+               if (err)
+                       return err;
+       }
+       if (en->en_result) {
+               err = visit(en->en_result, arg);
+               if (err)
+                       return err;
+       }
        return 0;
 }
 
 static PyObject *
 enum_next(enumobject *en)
 {
-       PyObject *result;
        PyObject *next_index;
        PyObject *next_item;
+       PyObject *result = en->en_result;
+       PyObject *it = en->en_sit;
 
-       result = PyTuple_New(2);
-       if (result == NULL)
+       next_item = (*it->ob_type->tp_iternext)(it);
+       if (next_item == NULL)
                return NULL;
 
        next_index = PyInt_FromLong(en->en_index);
        if (next_index == NULL) {
-               Py_DECREF(result);
+               Py_DECREF(next_item);
                return NULL;
        }
-       PyTuple_SET_ITEM(result, 0, next_index);
+       en->en_index++;
 
-       next_item = PyIter_Next(en->en_sit);
-       if (next_item == NULL) {
-               Py_DECREF(result);
-               return NULL;
+       if (result->ob_refcnt == 1) {
+               Py_INCREF(result);
+               Py_DECREF(PyTuple_GET_ITEM(result, 0));
+               Py_DECREF(PyTuple_GET_ITEM(result, 1));
+       } else {
+               result = PyTuple_New(2);
+               if (result == NULL) {
+                       Py_DECREF(next_index);
+                       Py_DECREF(next_item);
+                       return NULL;
+               }
        }
-
-       en->en_index++;
+       PyTuple_SET_ITEM(result, 0, next_index);
        PyTuple_SET_ITEM(result, 1, next_item);
        return result;
 }