File: tee.py

package info (click to toggle)
python-picklable-itertools 0.1.1-2
  • links: PTS, VCS
  • area: main
  • in suites: buster
  • size: 172 kB
  • sloc: python: 1,222; makefile: 3
file content (45 lines) | stat: -rw-r--r-- 1,328 bytes parent folder | download | duplicates (3)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
"""Support code for implementing `tee`."""
import collections
import six
from picklable_itertools import iter_


class tee_iterator(six.Iterator):
    """An iterator that works in conjunction with a `tee_manager`."""
    def __init__(self, deque, manager):
        self._deque = deque
        self._manager = manager

    def __iter__(self):
        return self

    def __next__(self):
        if len(self._deque) > 0:
            return self._deque.popleft()
        else:
            self._manager.advance()
            assert len(self._deque) > 0
            return next(self)


class tee_manager(object):
    """An object that manages a base iterator and publishes results to
    one or more client `tee_iterators`.
    """
    def __init__(self, iterable, n=2):
        self._iterable = iter_(iterable)
        self._deques = tuple(collections.deque() for _ in range(n))

    def iterators(self):
        return tuple(tee_iterator(deque, self) for deque in self._deques)

    def advance(self):
        """Advance the base iterator, publish to constituent iterators."""
        elem = next(self._iterable)
        for deque in self._deques:
            deque.append(elem)


def tee(iterable, n=2):
    """tee(iterable, n=2) --> tuple of n independent iterators."""
    return tee_manager(iter_(iterable), n=n).iterators()