Объект в цикле for: __iter__

for p in menu.items работает, но открывает внутренности меню. Хочется писать просто for p in menu: — и чтобы так же работали list(menu), sum(...) и генераторы списков. Для этого объекту нужен метод __iter__.

__iter__ и yield

for x in obj начинается с вызова obj.__iter__(). Проще всего написать этот метод с yield: функция или метод, в теле которого есть yield, — это генератор (generator). Он не возвращает значение один раз, как return, а выдаёт значения по одному: yield отдаёт очередное значение циклу и замирает до следующего шага. Когда тело метода закончилось, цикл тоже заканчивается.

def countdown(n):
    while n > 0:
        yield n  # выдать значение и замереть до следующего шага
        n -= 1


for x in countdown(3):
    print(x)
print(list(countdown(5)))

3
2
1
[5, 4, 3, 2, 1]

В __iter__ можно и не писать yield: достаточно вернуть итератор по внутреннему списку — iter(…) (о нём подробнее ниже). Важно одно — for получает новый итератор при каждом обходе:

class Shift:
    def __init__(self, barista, hours):
        self.barista = barista
        self._hours = list(hours)

    def __iter__(self):
        return iter(self._hours)  # новый итератор по списку часов


morning = Shift("Анна", ["08:00", "09:00", "10:00"])
for hour in morning:
    print(morning.barista, hour)
print(list(morning), len(list(morning)))

Анна 08:00
Анна 09:00
Анна 10:00
['08:00', '09:00', '10:00'] 3

class Product:
    def __init__(self, name, price):
        self.name = name
        self.price = price

    def __repr__(self):
        return f"Product({self.name!r}, {self.price})"


class Menu:
    def __init__(self, products):
        self._items = list(products)

    def __iter__(self):
        ...

class Product:
    def __init__(self, name, price):
        self.name = name
        self.price = price

    def __repr__(self):
        return f"Product({self.name!r}, {self.price})"


class Menu:
    def __init__(self, products):
        self._items = list(products)

    def __iter__(self):
        for p in self._items:
            yield p

Генератор в обычном методе

yield можно писать и в обычном методе — тогда метод возвращает генератор, по которому можно пройти циклом. Это удобно для «выборок» без промежуточных списков:

class Menu:
    # скопируйте сюда __init__ и __iter__ из упражнения выше

    def cheaper_than(self, price):
        ...

class Menu:
    def __init__(self, products):
        self._items = list(products)

    def __iter__(self):
        for p in self._items:
            yield p

    def cheaper_than(self, price):
        for p in self._items:
            if p.price < price:
                yield p

Итерируемый объект и итератор

За for стоят два понятия. Итерируемый объект (iterable) — то, по чему можно пройти: список, строка, меню с __iter__. Итератор (iterator) — объект, который выдаёт элементы по одному и помнит, где остановился. iter(x) просит у итерируемого объекта новый итератор, а next(it) берёт из итератора следующий элемент. Когда элементы кончились, next выбрасывает StopIteration — так for узнаёт, что пора остановиться:

drinks = ["латте", "чай"]
it = iter(drinks)
print(next(it))
print(next(it))
try:
    next(it)
except StopIteration:
    print("StopIteration — элементы кончились")

латте
чай
StopIteration — элементы кончились

Главное отличие: список, строку или меню с __iter__ можно проходить сколько угодно раз — каждый for вызывает __iter__ и получает новый итератор. Итератор тоже можно пройти циклом for (iter(it) возвращает его самого), но только один раз: дошёл до конца — пуст навсегда. Генератор — тоже итератор: результат menu.cheaper_than(200) из упражнения выше хватает на один проход, для второго вызовите метод снова.

it = iter(drinks)
print(list(it))
print(list(it))  # итератор уже пройден
print(list(drinks), list(drinks))  # список — сколько угодно раз

['латте', 'чай']
[]
['латте', 'чай'] ['латте', 'чай']

class Playlist:
    def __init__(self, songs):
        self._it = iter(songs)

    def __iter__(self):
        return self._it

class Playlist:
    def __init__(self, songs):
        self._songs = list(songs)

    def __iter__(self):
        for song in self._songs:
            yield song