* 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 701b0c0..00f97ef 100644
--- 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 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 @@
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 66ee870..5fbecbf 100644
--- a/Objects/enumobject.c
+++ b/Objects/enumobject.c
@@ -6,6 +6,7 @@
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 @@
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 @@
{
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;
}
+ en->en_index++;
+
+ 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;
+ }
+ }
PyTuple_SET_ITEM(result, 0, next_index);
-
- next_item = PyIter_Next(en->en_sit);
- if (next_item == NULL) {
- Py_DECREF(result);
- return NULL;
- }
-
- en->en_index++;
PyTuple_SET_ITEM(result, 1, next_item);
return result;
}