Author: Armin Rigo <ar...@tunes.org>
Branch:
Changeset: r85594:bca9142297fe
Date: 2016-07-07 09:00 +0200
http://bitbucket.org/pypy/pypy/changeset/bca9142297fe/
Log: Issue #2340: itertools.tee(x) uses a different logic if
hasattr(iter(x), '__copy__').
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
@@ -835,7 +835,7 @@
data before the other iterator, it is faster to use list() instead
of tee()
- Equivalent to :
+ If iter(iterable) has no method __copy__(), this is equivalent to:
def tee(iterable, n=2):
def gen(next, data={}, cnt=[0]):
@@ -848,17 +848,22 @@
yield item
it = iter(iterable)
return tuple([gen(it.next) for i in range(n)])
+
+ If iter(iterable) has a __copy__ method, though, we just return
+ a tuple t = (iterable, t[0].__copy__(), t[1].__copy__(), ...).
"""
if n < 0:
raise oefmt(space.w_ValueError, "n must be >= 0")
- if isinstance(w_iterable, W_TeeIterable): # optimization only
- chained_list = w_iterable.chained_list
- w_iterator = w_iterable.w_iterator
+ if space.findattr(w_iterable, space.wrap("__copy__")) is not None:
+ # In this case, we don't instantiate any W_TeeIterable.
+ # We just rely on doing repeated __copy__(). This case
+ # includes the situation where w_iterable is already
+ # a W_TeeIterable itself.
iterators_w = [w_iterable] * n
for i in range(1, n):
- iterators_w[i] = space.wrap(W_TeeIterable(space, w_iterator,
- chained_list))
+ w_iterable = space.call_method(w_iterable, "__copy__")
+ iterators_w[i] = w_iterable
else:
w_iterator = space.iter(w_iterable)
chained_list = TeeChainedListNode()
@@ -891,6 +896,11 @@
self.chained_list = chained_list.next
return w_obj
+ def copy_w(self):
+ space = self.space
+ tee_iter = W_TeeIterable(space, self.w_iterator, self.chained_list)
+ return space.wrap(tee_iter)
+
def W_TeeIterable___new__(space, w_subtype, w_iterable):
# Obscure and undocumented function. PyPy only supports w_iterable
# being a W_TeeIterable, because the case where it is a general
@@ -906,6 +916,7 @@
__new__ = interp2app(W_TeeIterable___new__),
__iter__ = interp2app(W_TeeIterable.iter_w),
next = interp2app(W_TeeIterable.next_w),
+ __copy__ = interp2app(W_TeeIterable.copy_w),
__weakref__ = make_weakref_descr(W_TeeIterable),
)
W_TeeIterable.typedef.acceptable_as_base_class = False
diff --git a/pypy/module/itertools/test/test_itertools.py
b/pypy/module/itertools/test/test_itertools.py
--- a/pypy/module/itertools/test/test_itertools.py
+++ b/pypy/module/itertools/test/test_itertools.py
@@ -702,6 +702,48 @@
x = d.next()
assert x == 'b'
+ def test_tee_defines_copy(self):
+ import itertools
+ a, b = itertools.tee('abc')
+ c = b.__copy__()
+ assert list(a) == ['a', 'b', 'c']
+ assert list(b) == ['a', 'b', 'c']
+ assert list(c) == ['a', 'b', 'c']
+ a, = itertools.tee('abc', 1)
+ x = a.next()
+ assert x == 'a'
+ b = a.__copy__()
+ x = a.next()
+ assert x == 'b'
+ x = b.next()
+ assert x == 'b'
+
+ def test_tee_function_uses_copy(self):
+ import itertools
+ class MyIterator(object):
+ def __iter__(self):
+ return self
+ def next(self):
+ raise NotImplementedError
+ def __copy__(self):
+ return iter('def')
+ my = MyIterator()
+ a, = itertools.tee(my, 1)
+ assert a is my
+ a, b = itertools.tee(my)
+ assert a is my
+ assert b is not my
+ assert list(b) == ['d', 'e', 'f']
+ # this gives AttributeError because it tries to call
+ # my.__copy__().__copy__() and there isn't one
+ raises(AttributeError, itertools.tee, my, 3)
+
+ def test_tee_function_empty(self):
+ import itertools
+ assert itertools.tee('abc', 0) == ()
+ a, = itertools.tee('abc', 1)
+ assert itertools.tee(a, 0) == ()
+
class AppTestItertools26:
spaceconfig = dict(usemodules=['itertools'])
_______________________________________________
pypy-commit mailing list
pypy-commit@python.org
https://mail.python.org/mailman/listinfo/pypy-commit
