Свой контейнер: __len__, __getitem__, __contains__

Корзина и меню из прошлых модулей хранят товары внутри, но снаружи с ними неудобно: чтобы узнать число товаров, пишут len(cart.items), а чтобы найти товар — вызывают метод. Со списком или словарём всё проще: len(menu), menu[0], "Латте" in menu. Такие объекты называют контейнерами. Свой класс тоже может стать контейнером — нужно объявить несколько магических методов. Набор методов, которые ждёт встроенная операция, называют протоколом (protocol): здесь — протоколом контейнера.

len и bool

len(x) вызывает x.__len__() — метод должен вернуть целое число не меньше нуля. А в if x: и bool(x) Python спрашивает x.__bool__(); если такого метода нет, но есть __len__, — объект ложен, когда его длина 0. Поэтому пустая корзина с __len__ сама собой становится «ложью»:

class Ticket:
    def __init__(self, number, lines):
        self.number = number
        self.lines = list(lines)

    def __len__(self):
        return len(self.lines)


full = Ticket(17, ["латте", "круассан"])
empty = Ticket(18, [])
print(len(full), len(empty))
for ticket in [full, empty]:
    if ticket:
        print(f"заказ №{ticket.number}: {len(ticket)} поз.")
    else:
        print(f"заказ №{ticket.number} пуст")

2 0
заказ №17: 2 поз.
заказ №18 пуст

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 Cart:
    def __init__(self):
        self.items = []

    def add(self, product):
        self.items.append(product)

    def __len__(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 Cart:
    def __init__(self):
        self.items = []

    def add(self, product):
        self.items.append(product)

    def __len__(self):
        return len(self.items)

__bool__ объявляют, когда «истинность» объекта означает не «непустой», а что-то другое:

class CoffeeMachine:
    def __init__(self, water, beans):
        self.water = water
        self.beans = beans

    def __bool__(self):
        ...

class CoffeeMachine:
    def __init__(self, water, beans):
        self.water = water
        self.beans = beans

    def __bool__(self):
        return self.water > 0 and self.beans > 0

Квадратные скобки: __getitem__

x[key] вызывает x.__getitem__(key). Что считать ключом, решает класс: список принимает номер, словарь — ключ. Например, расписание смен удобно открывать по дню недели:

class Schedule:
    def __init__(self, shifts):
        self._shifts = dict(shifts)  # день → бариста

    def __getitem__(self, day):
        if day not in self._shifts:
            raise KeyError(day)
        return self._shifts[day]


week = Schedule({"пн": "Анна", "вт": "Борис", "ср": "Анна"})
print(week["вт"])
try:
    week["вс"]
except KeyError as error:
    print("нет смены:", error)

Борис
нет смены: 'вс'

Если такого ключа нет, по соглашению выбрасывают KeyError (для номеров — IndexError), как делает словарь.

in: __contains__

x in container вызывает container.__contains__(x). Без него Python перебрал бы элементы по одному — но только если объект умеет перебираться. С __contains__ вы сами решаете, что значит «есть в контейнере»:

class OpenDays:
    def __init__(self, days):
        self._days = list(days)

    def __contains__(self, day):
        return day in self._days


open_days = OpenDays(["пн", "вт", "ср", "чт", "пт", "сб"])
print("сб" in open_days, "вс" in open_days)

True False

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

    def __getitem__(self, name):
        ...

    def __contains__(self, name):
        ...

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

    def __getitem__(self, name):
        for p in self._items:
            if p.name == name:
                return p
        raise KeyError(name)

    def __contains__(self, name):
        for p in self._items:
            if p.name == name:
                return True
        return False