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