My implementation of itertools.tee function

my_tee.py

from itertools import tee


def source():
    for i in range(5):
        print(f'next is {i}')
        yield i



def my_tee(it, n=2):
    # храним в памяти кэш из всех полученных значений
    values = []

    # новый итератор с кэшированием
    def _internal_iterator():
        # счетчик текущего элемента (свой для каждой копии итератора)
        i = 0
        try:
            # мы не знаем длину, поэтому цилк бесконечен до StopIteration
            while True:
                # если этого элемента нет еще в кэше (мы первые обратились за ним)
                if i >= len(values):
                    # извлечем его и сохраним, next может бросить StopIteration, если исходный итератор исчерпан
                    item = next(it)
                    values.append(item)
                else:
                    # иначе кто-то уже обращался за i-тым элементом и сохранил его
                    item = values[i]
                # отдадим элемент и сдвинем счетчик
                yield item
                i += 1
        except StopIteration:
            # цикл прерывается только StopIteration из исходного итератора
            pass

    # сделаем n копий нового итератора
    return [_internal_iterator() for _ in range(n)]


def main():
    coffee, tea, me = my_tee(source(), 3)

    print('get coffee = ', next(coffee))
    print('get coffee = ', next(coffee))
    print('get coffee = ', next(coffee))

    for i in tea:
        print('get tea = ', i)

    for i in me:
        print('get me = ', i)


main()