tee.py

import itertools
import random


def mytee1(xs, n):
    last = [None, None]

    def gen(it, mylast):
        nonlocal last
        while True:
            mylast = mylast[1]
            if not mylast:
                mylast = last[1] = last = [next(it), None]
            yield mylast[0]

    it = iter(xs)
    return tuple(gen(it, last) for _ in range(n))


def run(func):
    print("===", func, "===")
    it = iter(range(15))

    niters = 4
    iters = func(it, niters)
    lists = [[] for i in range(niters)]
    stopped = [False] * niters

    rnd = random.Random(42)
    while True:
        i = rnd.randrange(niters)
        if stopped[i]:
            continue
        try:
            lists[i].append(next(iters[i]))
        except StopIteration:
            stopped[i] = True
            if all(stopped):
                break

    for i in range(niters):
        print(lists[i])


run(itertools.tee)
run(mytee1)