Классы-итераторы и генераторы в Python
В языке программирования Python есть понятия итераторов и генераторов. Оба термина не очень сложные и часто является обычными вопросами для собеседований разработчиков. Но в данной статье хочется разобрать немного примеров как создавать подобные объекты с помощью классов, так как это достаточно мощный и удобный механизм.
Итераторы и генераторы
В определении на сайте doc.python.org говорится, что объект-итератор должен
- Содержать метод
__next__()
, вызываем встроенным методомnext()
для возвращения следующего элемента коллекции. - Если данных больше нет, то должно возникать исключение
StopIteration
. - Также обязателен метод
__iter__()
для возвращения объекта-итератора.
Примеры встроенных типов, примеров итераторов это list
, tuple
.
С генераторами тоже все довольно просто. Это объекты, возвращающие функцию генерации с выражением yield
. Её особенность в том, что при следующем вызове, функция продолжает свою работу не с самого начала, а с места использования этого служебного слова.
from typing import Generator
def generator(start: int, stop: int, step: int = 1) -> Generator[int, None, None]:
"""
Генератор целых чисел от start до stop с шагом step.
>>> list(generator(5, 13, 2))
[5, 7, 9, 11, 13]
"""
i = start
while i <= stop:
yield i
i += step
Особенно хочется подчеркнуть, что существуют удобные выражения-генераторы, когда отдельная функция не нужна.
Классы
Посмотрим на примеры классов, которые являются итераторами. Например в стандартной библиотеке есть пакет csv
, создав объект reader, можно итерироваться непосредственно по нему:
import csv
csv_file = open('my_file.csv')
my_reader = csv.reader(csv_file)
for row in my_reader:
print(row)
Кроме того в данном примере файл не парсится весь сразу, а данные возвращаются по ходу его чтения. Это уже поведения генератора.
Попытается написать класс с похожим поведением.
Простой генератор
Пример класса, который может выступать генератором, возвращающим квадраты чисел из входного списка.
class SquareIterator:
"""
R квадратов чисел.
>>> s = SquareIterator([3, 5, 7])
>>> list(s)
[9, 25, 49]
>>> list(s)
[]
"""
def __init__(self, data: List[int]) -> None:
self.index = 0
self.total = len(data)
self.data = data
def __iter__(self) -> Iterator[int]:
return self
def __next__(self) -> int:
if self.index >= self.total:
raise StopIteration
value = self.data[self.index] ** 2
self.index += 1
return value
Но иногда хочется иметь возможность итерироваться по коллекции несколько раз, как например мы можем это делать по типу list
. Для этого нужно дополнить метод __iter__
сбросом начального состояния.
class SquareInfIterator(SquareIterator):
"""
Генератор квадратов чисел с возможностью повторного вызова.
>>> s = SquareIterator([3, 5, 7])
>>> list(s)
[9, 25, 49]
>>> list(s)
[9б 25 49]
"""
def __iter__(self) -> Iterator[int]:
self.index = 0
return self
Произвольный генератор
Метод __iter__
не обязан всегда возвращать self
. В качестве возвращаемого значения ожидается любой генератор, который потом можно использовать в качестве аргумента для next()
:
class RangeGenerator:
"""
Range генератор.
>>> r = RangeGenerator(5, 13, 2)
>>> list(r)
[5, 7, 9, 11, 13]
"""
def __init__(self, start: int, stop: int, step: int = 1) -> None:
self.start = start
self.stop = stop
self.step = step
def __iter__(self) -> Generator[int, None, None]:
i = self.start
while i <= self.stop:
yield i
i += self.step
Аналогичным образом можно написать генератор, похожий на csv.reader
, когда не нужно читать файл целиком и нам интересны операции на отдельными строками:
class FileLineLenGenerator:
"""
Генератор длин строк файла.
>>> f = FileLineLenGenerator('my_file.txt')
>>> list(f)
[4, 3, 8]
"""
def __init__(self, file_name: str) -> None:
self.file_name = file_name
def __iter__(self) -> Generator[int, None, None]:
with open(self.file_name) as f:
for line in f:
yield len(line)
Заключение
Python итераторы это очень мощный инструмент языка. С помощью генераторов можно совершать "ленивые" вычисления, экономив на памяти. А возможность создавать произвольные классы итераторы делает программирование более гибким.
Все примеры доступны на github.com.