Coffee, tee or me

Итератор, как известно, выдает значения по одному (например, методом next), и его нельзя «отмотать» назад. Это означает, что получать все подряд значения из итератора может только один потребитель. Однако, если несколько потребителей хотят читать из одного итератора, его можно разделить с помощью функции itertools.tee. Кстати, tee переводится как тройник (тройник похож на букву Т).

tee принимает исходный итератор и число – количество новых итераторов, на которые разделится исходный, а возвращает кортеж из новых итераторов. При этом извлечение значений из одного из итераторов не влияет на остальные.

📎 Пример:

from itertools import tee

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

# три потребителя
coffee, tea, me = tee(source(), 3)

# первый берет два числа
next(coffee); next(coffee)
# второй одно
next(tea)
# третий - все
for i in me:
    ...

Вывод:

next is 0
next is 1
next is 2
next is 3
next is 4

Видно, что все значения были извлечены из исходного оператора по одному разу, но каждый потребитель получил их столько раз, сколько хотел.

Предостережения:

1. Исходный итератор не следует итерировать, иначе производные итераторы лишатся некоторых значений. 

2. Механизм tee таков, что он хранит в памяти извлеченные элементы, чтобы остальные потребители могли их получить, даже если исходный итератор уже сместился. Поэтому, если элементов много или они большие, это может серьезно повлиять на расход памяти.

Полный пример кода (плюс моя реализация tee на скорую руку) – по ссылке.

🐉 Специально для канала @pyway. Подписывайтесь на мой канал в Телеграм @pyway 👈 

Итераторы и генераторы

В чем разница между итератором и генератором? Этот вопрос можно часто услышать на собеседованиях.

Итератор – более общая концепция, чем генератор.

Итератор – это интерфейс доступа к элементам коллекций и потоков данных. Он требует реализации единственного метода – «дай мне следующий элемент». Если вы пишите свой итератор на Python 3 вам нужно реализовать в классе метод __next__. Если элементы исчерпаны итератор возбудит исключение StopIteration.

📎 Пример. Итератор счетчик – выдает числа от low до high:

class Counter:
    def __init__(self, low, high):
        self.current = low
        self.high = high
    def __iter__(self):
        return self
    def __next__(self): 
        if self.current > self.high:
            raise StopIteration
        else:
            self.current += 1
            return self.current - 1

Генератор – это итератор

Генератор – это итератор, но не наоборот. Не любой итератор является генератором.

Есть два способа получить генератор:

📎 1. Генераторное выражение (что-то типа list comprehension, но возвращает генератор, а не список). Используются круглые скобки:

>>> g = (2 * i for i in range(5))
>>> type(g)
<class 'generator'>
>>> next(g)
0
>>> next(g)
2

📎 2. Генераторные функции. Это функции, где есть хотя бы одно выражение yield. Когда мы запускаем генератор, функция выполняет до первого выражения yield. То, что мы передали в yield будет возвращено наружу. Генератор при этом встанет «на паузу» до следующей итерации. При следующей итерации выполнение генератора продолжится до очередного yield.

Генераторы можно прочитать только 1 раз, потому что обычно генераторы не хранят значения в памяти, а генерируют их налету (отсюда и название).

Пример. Генератор чисел Фибоначчи (бесконечный):

def fib():
    a, b = 0, 1
    while 1:
        yield a
        a, b = b, a + b

>>> fib_g = fib()
>>> next(fib_g)
0
>>> next(fib_g)
1
>>> next(fib_g)
1
>>> next(fib_g)
2
>>> next(fib_g)
3
>>> next(fib_g)
5

Вызвав генераторную функцию fib() мы получили генератор. Затем мы итерируем этот генератор функцией next().

Остановка генератора

Если генератор «закончился» (т.е. просто вышли из функции генератора в конце его кода или по return), то автоматически возбуждается исключение StopIteration. Это не ошибка, это нормально, просто принятый способ обработки конца итератора.

def gen():
    yield 1
    yield 5
    # и все, код кончился, вышли
    
for x in gen():
    print(x) # 1, 5

for in сам ловит исключение StopIteration и просто завершает итерировать этот генератор.

Передача данных в генератор

У генераторов есть дополнительные методы, которые позволяют передавать внутрь генератора данные или возбуждать внутри него исключения. Это еще одно отличие от простых итераторов.

send() – отправить данные в генератор. Переданное значение вернется из той конструкции yield, на которой возникла последняя пауза генератора. При этом генератор будет прокручен на один шаг, как если бы мы вызвали next:

val = yield i  # генератор вернет i, но внутри получит val из аргумента метода send

Пример. Этот генератор просто выдает числа от 0 и далее, при этом печатает в поток вывода все, что мы ему отправляем.

def my_gen():
    i = 0
    while True:
        val = yield i
        print('Got inside generator:', val)
        i += 1

>>> g = my_gen()
>>> next(g)
0
>>> g.send("hello")
Got inside generator: hello
1
>>> g.send("world")
Got inside generator: world
2

Обратите внимание, что первый раз нельзя посылать в генератор данные, пока мы не прокрутили его до первого yield. Нужно либо взывать next(g) или g.send(None) – это одно и тоже.

Не будет ошибкой отправлять данные генератору, который не получает их (нет использования значения конструкции yield). Например, нашему генератору fib() можно отравить все, что угодно, он просто проигнорирует.

throw() – бросить исключение внутри генератора. Исключение будет возбуждено из того выражение yield, где генератор последний раз остановился.

>>> g = my_gen()   # my_gen из прошлого примера

>>> g.throw(TypeError, 'my error')
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<stdin>", line 1, in my_gen
TypeError: my error

close() – закрыть генератор. Бросает внутри генератора особое исключение GeneratorExit. Это исключение, даже если оно не обработано, не распространится в код, вызвавший close(). Но, если мы поймали это исключение внутри генератора, то после закрытия генератора нельзя уже делать yield, рискуя получить RuntimeError. Остальные виды исключений будут распространяться из генератора в код, его вызывающий. Попытка итерировать закрытый итератор приведет к исключению StopIteration (закрытый генератор – пустой итератор).

>>> g = my_gen()
>>> next(g)
0
>>> next(g)
Got inside generator: None
1
>>> g.close()
>>> next(g)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
StopIteration

Бонус

Как взять из итератора (в том числе из генератора) N первых значений?

Можно, конечно, написать свою функцию. Но зачем, если она уже есть в стандартном модуле itertools. Этот модуль содержит множество вспомогательных функций для работы с итераторами. Нам понадобится itertools.islice. Первый аргумент – итератор (ну или генератор), остальные три – как в range.

>>> list(itertools.islice(fib(), 10))
[0, 1, 1, 2, 3, 5, 8, 13, 21, 34]

>>> list(itertools.islice(fib(), 10, 20, 2))
[55, 144, 377, 987, 2584]

В первом примере мы передаем в функцию itertools.islice наш генератор чисел Фибоначчи и число чисел, которые надо вычислить (в нашем случае – 10).

Мы также применяем функцию list, чтобы посмотреть список значений, потому что itertools.islice возвращает не спикок, а именно новый итератор, в котором будут только интересные нам значений из исходного итератора.

Во втором примеры аргументов 4 штуки. В этом случае второй аргумент – начальный номер = 10, третий – конечный номер = 20 – (не включительно), и четвертый – шаг = 2. (Очень похоже на range, не так ли?)

Специально для канала @pyway.