Задачка об итераторе на Python

На днях я предложил читателям телеграм-канала Oh My Py задачку об итераторе с ограничениями. Давайте её разберём.

Условия задачи

Допустим, вы основали модный HR-стартап, который подбирает идеальные коллективы сотрудников. Дело это нелёгкое, так что начали с простой эвристики:

Любой коллектив идеален, пока в нём не появляется Френк

Подготовили интеллектуальный алгоритм, который предлагает сотрудника:

import random
names = ["Френк", "Клер", "Зоя", "Питер", "Лукас"]

def employee():
  name = random.choice(names)
  return name

Остался последний шаг — разработать нечто под названием employeficator(), что и будет подбирать дружный коллектив. Использоваться будет так:

>>> [name for name in employeficator()]
['Зоя', 'Зоя', 'Питер']
>>> [name for name in employeficator()]
['Лукас', 'Зоя', 'Питер']

Ваша задача — реализовать employeficator() максимально идиоматично.

Для затравки начну заведомо неудачным вариантом:

def employeficator():
  employees = []
  name = employee()
  while name != "Френк":
    employees.append(name)
    name = employee()
  return employees

Решение-победитель

Для начала, что такое «идиоматично». Идиоматичный код использует «родные» конструкции языка и стандартной библиотеки, не нарушая при этом питонячий дзен (simple is better than complex, readability counts, вот это всё).

Месиво из вложенных циклов с break и continue вряд ли можно назвать идиоматичным. Точно также не будет идиоматичной «функциональная» колбаса из вызовов functools и itertools. Абсолютных критериев тут нет, но общий смысл, надеюсь, понятен.

Теперь к решению. Задача была с небольшим подвохом: искомый employeficator() уже есть в стандартной библиотеке. Больше того, не просто в стандартной библиотеке, а в самом её сердце, в built-in функциях! Вот он:

[name for name in iter(employee, "Френк")]

Да, это функция iter(). Обычно её вызывают с одним аргументом — коллекцией:

>>> seq = [1, 2, 3]
>>> it = iter(seq)
>>> next(it)
1

Но в варианте с двумя аргументами iter() работает иначе:

iter(callable, sentinel)

Первый аргумент — функция или что-нибудь вызываемое (callable), второй — контрольное значение (sentinel). Каждое обращение к итератору вызывает callable() и возвращает результат его выполнения. А как только callable() возвращает значение sentinel, итератор прекращает работу.

Это ровно то поведение, что требовалось в задаче — вызывать employee(), пока очередной вызов не вернёт "Френк".

Так что iter() здесь — идеальное решение. Поздравляю всех, кто его предложил!

Хорошие решения

Удачное решение — использовать генератор. Благо, в питоне 3.8 появилась короткая форма записи для инициализации переменной внутри выражения («моржовый» оператор):

def employeficator():
    while (name := employee()) != "Френк":
        yield name

«Морж» вызвал большое недовольство в питонячьем мире, так что если он вам не по душе, то можно и так:

def employeficator():
    name = employee()
    while name != "Френк":
        yield name
        name = employee()

Неудачные решения

Часто предлагали такой вариант:

def employeficator():
    while True:
        name = employee()
        if name == "Френк":
            break
        yield name

Ничего плохого в нём нет, но если есть возможность обойтись без break, не усложняя код — лучше это сделать (см. решение из предыдущего раздела).

Бывает, людям хочется применить «функциональщину»:

from itertools import takewhile, count

def employeficator():
    return takewhile(
        lambda name: name != 'Френк', (employee() for _ in count())
    )

from itertools import takewhile, starmap, repeat

def employeficator():
    return takewhile(
        lambda name: name != "Френк", starmap(employee, repeat(()))
    )

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

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

Хотя нет, одно всё же покажу:

def employeficator():
    team_size = range(0, random.randint(1, len(names) + 1))
    return iter(set(name if name != 'Френк' else
               next(employeficator()) for name in
               (employee() for _ in team_size)))

Пожалуйста, никогда не пишите так продакшен-код. Пожалейте коллег и себя.