Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit e8b0f04

Browse files
rhettingerrhettinger
committed
* Beefed-up tests
* Allow tuple re-use * Call tp_iternext directly
1 parent a95d3b7 commit e8b0f04

2 files changed

Lines changed: 87 additions & 27 deletions

File tree

‎Lib/test/test_enumerate.py‎

Lines changed: 45 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,8 @@
11
import unittest
2+
from sets import Set
23

34
from test import test_support
45

5-
seq, res = 'abc', [(0,'a'), (1,'b'), (2,'c')]
6-
76
class G:
87
'Sequence using __getitem__'
98
def __init__(self, seqn):
@@ -65,37 +64,49 @@ def __iter__(self):
6564
class EnumerateTestCase(unittest.TestCase):
6665

6766
enum = enumerate
67+
seq, res = 'abc', [(0,'a'), (1,'b'), (2,'c')]
6868

6969
def test_basicfunction(self):
70-
self.assertEqual(type(self.enum(seq)), self.enum)
71-
e = self.enum(seq)
70+
self.assertEqual(type(self.enum(self.seq)), self.enum)
71+
e = self.enum(self.seq)
7272
self.assertEqual(iter(e), e)
73-
self.assertEqual(list(self.enum(seq)), res)
73+
self.assertEqual(list(self.enum(self.seq)), self.res)
7474
self.enum.__doc__
7575

7676
def test_getitemseqn(self):
77-
self.assertEqual(list(self.enum(G(seq))), res)
77+
self.assertEqual(list(self.enum(G(self.seq))), self.res)
7878
e = self.enum(G(''))
7979
self.assertRaises(StopIteration, e.next)
8080

8181
def test_iteratorseqn(self):
82-
self.assertEqual(list(self.enum(I(seq))), res)
82+
self.assertEqual(list(self.enum(I(self.seq))), self.res)
8383
e = self.enum(I(''))
8484
self.assertRaises(StopIteration, e.next)
8585

8686
def test_iteratorgenerator(self):
87-
self.assertEqual(list(self.enum(Ig(seq))), res)
87+
self.assertEqual(list(self.enum(Ig(self.seq))), self.res)
8888
e = self.enum(Ig(''))
8989
self.assertRaises(StopIteration, e.next)
9090

9191
def test_noniterable(self):
92-
self.assertRaises(TypeError, self.enum, X(seq))
92+
self.assertRaises(TypeError, self.enum, X(self.seq))
9393

9494
def test_illformediterable(self):
95-
self.assertRaises(TypeError, list, self.enum(N(seq)))
95+
self.assertRaises(TypeError, list, self.enum(N(self.seq)))
9696

9797
def test_exception_propagation(self):
98-
self.assertRaises(ZeroDivisionError, list, self.enum(E(seq)))
98+
self.assertRaises(ZeroDivisionError, list, self.enum(E(self.seq)))
99+
100+
def test_argumentcheck(self):
101+
self.assertRaises(TypeError, self.enum) # no arguments
102+
self.assertRaises(TypeError, self.enum, 1) # wrong type (not iterable)
103+
self.assertRaises(TypeError, self.enum, 'abc', 2) # too many arguments
104+
105+
def test_tuple_reuse(self):
106+
# Tests an implementation detail where tuple is reused
107+
# whenever nothing else holds a reference to it
108+
self.assertEqual(len(Set(map(id, list(self.seq)))), len(self.seq))
109+
self.assertEqual(len(Set(map(id, enumerate(self.seq)))), min(1,len(self.seq)))
99110

100111
class MyEnum(enumerate):
101112
pass
@@ -104,8 +115,28 @@ class SubclassTestCase(EnumerateTestCase):
104115

105116
enum = MyEnum
106117

107-
def test_main():
108-
test_support.run_unittest(EnumerateTestCase, SubclassTestCase)
118+
class TestEmpty(EnumerateTestCase):
119+
120+
seq, res = '', []
121+
122+
class TestBig(EnumerateTestCase):
123+
124+
seq = range(10,20000,2)
125+
res = zip(range(20000), seq)
126+
127+
128+
def test_main(verbose=None):
129+
testclasses = (EnumerateTestCase, SubclassTestCase, TestEmpty, TestBig)
130+
test_support.run_unittest(*testclasses)
131+
132+
# verify reference counting
133+
import sys
134+
if verbose and hasattr(sys, "gettotalrefcount"):
135+
counts = [None] * 5
136+
for i in xrange(len(counts)):
137+
test_support.run_unittest(*testclasses)
138+
counts[i] = sys.gettotalrefcount()
139+
print counts
109140

110141
if __name__ == "__main__":
111-
test_main()
142+
test_main(verbose=True)

‎Objects/enumobject.c‎

Lines changed: 42 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ typedef struct {
66
PyObject_HEAD
77
long en_index; /* current index of enumeration */
88
PyObject* en_sit; /* secondary iterator of enumeration */
9+
PyObject* en_result; /* result tuple */
910
} enumobject;
1011

1112
PyTypeObject PyEnum_Type;
@@ -30,6 +31,16 @@ enum_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
3031
Py_DECREF(en);
3132
return NULL;
3233
}
34+
en->en_result = PyTuple_New(2);
35+
if (en->en_result == NULL) {
36+
Py_DECREF(en->en_sit);
37+
Py_DECREF(en);
38+
return NULL;
39+
}
40+
Py_INCREF(Py_None);
41+
PyTuple_SET_ITEM(en->en_result, 0, Py_None);
42+
Py_INCREF(Py_None);
43+
PyTuple_SET_ITEM(en->en_result, 1, Py_None);
3344
return (PyObject *)en;
3445
}
3546

@@ -38,42 +49,60 @@ enum_dealloc(enumobject *en)
3849
{
3950
PyObject_GC_UnTrack(en);
4051
Py_XDECREF(en->en_sit);
52+
Py_XDECREF(en->en_result);
4153
en->ob_type->tp_free(en);
4254
}
4355

4456
static int
4557
enum_traverse(enumobject *en, visitproc visit, void *arg)
4658
{
47-
if (en->en_sit)
48-
return visit(en->en_sit, arg);
59+
int err;
60+
61+
if (en->en_sit) {
62+
err = visit(en->en_sit, arg);
63+
if (err)
64+
return err;
65+
}
66+
if (en->en_result) {
67+
err = visit(en->en_result, arg);
68+
if (err)
69+
return err;
70+
}
4971
return 0;
5072
}
5173

5274
static PyObject *
5375
enum_next(enumobject *en)
5476
{
55-
PyObject *result;
5677
PyObject *next_index;
5778
PyObject *next_item;
79+
PyObject *result = en->en_result;
80+
PyObject *it = en->en_sit;
5881

59-
result = PyTuple_New(2);
60-
if (result == NULL)
82+
next_item = (*it->ob_type->tp_iternext)(it);
83+
if (next_item == NULL)
6184
return NULL;
6285

6386
next_index = PyInt_FromLong(en->en_index);
6487
if (next_index == NULL) {
65-
Py_DECREF(result);
88+
Py_DECREF(next_item);
6689
return NULL;
6790
}
68-
PyTuple_SET_ITEM(result, 0, next_index);
91+
en->en_index++;
6992

70-
next_item = PyIter_Next(en->en_sit);
71-
if (next_item == NULL) {
72-
Py_DECREF(result);
73-
return NULL;
93+
if (result->ob_refcnt == 1) {
94+
Py_INCREF(result);
95+
Py_DECREF(PyTuple_GET_ITEM(result, 0));
96+
Py_DECREF(PyTuple_GET_ITEM(result, 1));
97+
} else {
98+
result = PyTuple_New(2);
99+
if (result == NULL) {
100+
Py_DECREF(next_index);
101+
Py_DECREF(next_item);
102+
return NULL;
103+
}
74104
}
75-
76-
en->en_index++;
105+
PyTuple_SET_ITEM(result, 0, next_index);
77106
PyTuple_SET_ITEM(result, 1, next_item);
78107
return result;
79108
}

0 commit comments

Comments
 (0)