Skip to content

Instantly share code, notes, and snippets.

@tirinox tirinox/my_tee.py
Created Aug 2, 2019

Embed
What would you like to do?
My implementation of itertools.tee function
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()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.