Почти десять лет назад я показал краткое введение в менеджеры контекста (пункт 2 здесь) и думал, что стану активнее пользоваться такими менеджерами. Но вспомнил я о них только недавно, на фоне того, как много мне приходилось перенастраивать и очищать тестируемый код при опытах по параллелизму (код получался одновременно неприятным и некрасивым).

Посмотрите спецификацию PEP 343: там описано, что суть менеджеров контекста в следующем: «позволить вычленять в отдельные блоки стандартные варианты использования инструкций try/finally». Мне всегда казалось, что finally тяготеет к обработке исключений. Но это не столько обработка ошибок, сколько очистка. Конечно, вы должны быть в состоянии обеспечить качественную очистку в случае, если выброшено исключение, но её к тому же необходимо обеспечить, несмотря на то, что вы покидаете область видимости. Думаю, здесь мы слишком полагались на вызовы функций как на основную рабочую единицу, что отвлекало нас от области видимости как от более общей концепции. Эта тема особенно интересна в сравнении с временами жизни в Rust.

В настоящее время я убеждён, что самый простой и красивый способ написать менеджер контекста – это воспользоваться декоратором функции @contextmanager, написав функцию как генератор. Но генератор – это функция с yield, и во что же она результирует? Ранее я уже приводил пример, в котором результатом (yield) был каталог, передававшийся в функцию. Но это фактически не имеет смысла, потому что (A) я так и не буду его использовать и (B) эта информация мне уже известна. Поэтому думаю, что хорошо было бы сделать так:

@contextmanager
def visitDir(d):
    old = os.getcwd()
    os.chdir(d)
    yield # No 'd'
    os.chdir(old)

Таким образом, ничего не будем yield. Этот код передаёт управление обратно создателю контекста, чтобы функции могли нормально работать, а когда область видимости закрывается, управление возвращается менеджеру контекста, и он выполняет os.chdir(old).
Но, всё-таки, что мы будем yield? Именно здесь нам пригодится ключевое слово as (добавленное для управления контекстом вместе с with). Если у вас есть менеджер контекста cm, и вы создаёте контекст вот так:

with cm() as x:
    ...

то результатом работы менеджера контекста является x, и он доступен на всём оставшемся протяжении области видимости.

В качестве такого результата может выдаваться любой объект. В моих примерах с конкурентностью мне приходилось передавать информацию в область видимости и возвращать оттуда информацию, поэтому я создал в scenario_tester.py новый тип, который назвал Scenario:

from contextlib import contextmanager
from dataclasses import dataclass, field
import time
import os
from pprint import pformat

@dataclass
class Scenario:
    multiplier: int
    tasks: int
    args1: range = field(init=False)
    args2: list[int] = field(init=False)
    results: list[float] = field(default_factory=list)

    def __post_init__(self):
        self.args1 = range(self.tasks)
        self.args2 = [self.multiplier] * self.tasks

@contextmanager
def scenario():
    multiplier = 1  # Increase for longer computations
    logical_processors = os.cpu_count()
    print(f"{logical_processors = }")
    tasks = (logical_processors - 0) * 1  # Try different numbers
    print(f"{tasks = }")
    start = time.monotonic()
    scenario = Scenario(multiplier, tasks)
    try:
        yield scenario
    finally:
        elapsed = time.monotonic() - start
        print(
            f"""{pformat(list(scenario.results))}
              Elapsed time: {elapsed:.2f}s"""
        )

Scenario создаёт и предоставляет args1 и args2 и подхватывает results от теста. Здесь я использую dataclass, поскольку теперь считаю его вариантом, используемым по умолчанию. Подробнее об этом рассказано в моей презентации с конференции Pycon за 2022 год. Далее __post_init__() создаёт args1 и args2, которые намеренно деинициализируются конструктором, сгенерированным через dataclass.

Функция менеджера контекста scenario() устанавливает всё необходимое, а затем создаёт объект Scenario и получает от него результат, который затем используется внутри контекста. Когда контекстная область видимости заканчивается, объект Scenario остаётся доступен, поэтому можно извлечь и отобразить scenario.results. Обратите внимание, что в сценарий не включено время start, поскольку внутри контекста оно не нужно, но всё равно может использоваться в finally, так как находится в области видимости у функции менеджера контекста.

Менеджер контекста позволяет избавиться от всего дополнительного кода, который я написал для каждого из тестов. Например, вот как выглядит ситуация с with_processes.py:

from concurrent.futures import ProcessPoolExecutor
from scenario_tester import scenario
from cpu_intensive import cpu_intensive

if __name__ == "__main__":
    with scenario() as scenario:
        with ProcessPoolExecutor() as executor:
            scenario.results = executor.map(
                cpu_intensive, scenario.args1, scenario.args2
            )

Выглядит почти идентично with_threads.py:

from concurrent.futures import ThreadPoolExecutor
from scenario_tester import scenario
from cpu_intensive import cpu_intensive

if __name__ == "__main__":
    with scenario() as scenario:
        with ThreadPoolExecutor() as executor:
            scenario.results = executor.map(
                cpu_intensive, scenario.args1, scenario.args2
            )

Другие примеры тоже похожи.

Отметив оставшийся дублирующийся код, можем пойти ещё дальше и передать тип исполнителя к функции как параметр, так, как это сделано в function_tester.py:

from concurrent.futures import ProcessPoolExecutor
from scenario_tester import scenario
from cpu_intensive import cpu_intensive

def test_cpu_intensive(ExecutorClass):
    with scenario() as s:
        with ExecutorClass() as executor:
            s.results = executor.map(cpu_intensive, s.args1, s.args2)

if __name__ == "__main__":
    test_cpu_intensive(ProcessPoolExecutor)

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