Author: Armin Rigo <[email protected]>
Branch:
Changeset: r57092:6bcaedad2d93
Date: 2012-09-03 11:52 +0200
http://bitbucket.org/pypy/pypy/changeset/6bcaedad2d93/
Log: Rewrite itertools.tee() to use a simple chained list internally,
instead of a list that only ever grows.
diff --git a/pypy/module/itertools/interp_itertools.py
b/pypy/module/itertools/interp_itertools.py
--- a/pypy/module/itertools/interp_itertools.py
+++ b/pypy/module/itertools/interp_itertools.py
@@ -811,7 +811,7 @@
item = data[i] = next()
cnt[0] += 1
else:
- item = data.pop(i)
+ item = data[i] # data.pop(i) if it's the last one
yield item
it = iter(iterable)
return tuple([gen(it.next) for i in range(n)])
@@ -821,47 +821,42 @@
myiter = space.interpclass_w(w_iterable)
if isinstance(myiter, W_TeeIterable): # optimization only
- tee_state = myiter.tee_state
+ chained_list = myiter.chained_list
+ w_iterator = myiter.w_iterator
iterators_w = [w_iterable] * n
for i in range(1, n):
- iterators_w[i] = space.wrap(W_TeeIterable(space, tee_state,
- myiter.index))
+ iterators_w[i] = space.wrap(W_TeeIterable(space, w_iterator,
+ chained_list))
else:
- tee_state = TeeState(space, w_iterable)
- iterators_w = [space.wrap(W_TeeIterable(space, tee_state)) for x in
range(n)]
+ w_iterator = space.iter(w_iterable)
+ chained_list = TeeChainedListNode()
+ iterators_w = [space.wrap(
+ W_TeeIterable(space, w_iterator, chained_list))
+ for x in range(n)]
return space.newtuple(iterators_w)
-class TeeState(object):
- def __init__(self, space, w_iterable):
- self.space = space
- self.w_iterable = self.space.iter(w_iterable)
- self.num_saved = 0
- self.saved_w = []
-
- def get_next(self, index):
- if index >= self.num_saved:
- w_obj = self.space.next(self.w_iterable)
- self.saved_w.append(w_obj)
- self.num_saved += 1
- return w_obj
- else:
- return self.saved_w[index]
+class TeeChainedListNode(object):
+ w_obj = None
class W_TeeIterable(Wrappable):
- def __init__(self, space, tee_state, start_index=0):
+ def __init__(self, space, w_iterator, chained_list):
self.space = space
- self.tee_state = tee_state
- self.index = start_index
+ self.w_iterator = w_iterator
+ assert chained_list is not None
+ self.chained_list = chained_list
def iter_w(self):
return self.space.wrap(self)
def next_w(self):
- try:
- w_obj = self.tee_state.get_next(self.index)
- return w_obj
- finally:
- self.index += 1
+ chained_list = self.chained_list
+ w_obj = chained_list.w_obj
+ if w_obj is None:
+ w_obj = self.space.next(self.w_iterator)
+ chained_list.next = TeeChainedListNode()
+ chained_list.w_obj = w_obj
+ self.chained_list = chained_list.next
+ return w_obj
def W_TeeIterable___new__(space, w_subtype, w_iterable):
# Obscure and undocumented function. PyPy only supports w_iterable
@@ -870,8 +865,8 @@
# semantics are then slightly different; see the XXX in lib-python's
# test_itertools).
myiter = space.interp_w(W_TeeIterable, w_iterable)
- tee_state = myiter.tee_state
- return space.wrap(W_TeeIterable(space, tee_state))
+ return space.wrap(W_TeeIterable(space, myiter.w_iterator,
+ myiter.chained_list))
W_TeeIterable.typedef = TypeDef(
'_tee',
_______________________________________________
pypy-commit mailing list
[email protected]
http://mail.python.org/mailman/listinfo/pypy-commit
