Условные или дополнительные контекстные менеджеры с оператором

Question

Предположим, у меня есть какой-то менеджер контекста (из сторонней библиотеки), который я использую следующим образом:

with freeze_time(test_dt):
    lines_of_code_1
    lines_of_code_2
    lines_of_code_3

Но предположим, что если для test_dt нет значения, менеджер контекста не должен запускаться, но должен выполняться весь оставшийся код, например:

if test_dt:
    with freeze_time(test_dt):
        lines_of_code_1
        lines_of_code_2
        lines_of_code_3
else:
    lines_of_code_1
    lines_of_code_2
    lines_of_code_3

Предположим, что lines_of_code здесь 2-3 строки кода, которые абсолютно идентичны, есть ли более чистый способ написания этого? Я знаю, что я мог бы написать что-то вроде этого:

def do_thing():
    lines_of_code_1
    lines_of_code_2
    lines_of_code_3

if test_dt:
    with freeze_time(test_dt):
        do_thing()
else:
    do_thing()

Но я не без ума от этого форматирования. Кроме того, я не хочу засорять этот шаблон по всему моему коду.

Есть одна последняя возможность, но я не уверен, что это сработает: создание подкласса менеджера контекста и пропуск __enter__ а также __exit__ функционирует, если заданный test_dt пуст, например:

class optional_freeze_time(object):
    def __init__(self, test_dt=None):
        if test_dt:
            self.ctx_manager = freeze_time(test_dt)
        else:
            self.ctx_manager = None
    def __enter__(self, *args, **kwargs):
        if self.ctx_manager:
            self.ctx_manager.__enter__(*args, **kwargs)
    def __exit__(self, *args, **kwargs):
        if self.ctx_manager:
            self.ctx_manager.__exit__(*args, **kwargs)

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

Answer 1

Вот простой способ обернуть существующий менеджер контекста, даже не используя классы:

from contextlib import contextmanager

@contextmanager
def example_context_manager():
    print('before')
    yield
    print('after')

@contextmanager
def optional(condition, context_manager):
    if condition:
        with context_manager:
            yield
    else:
        yield

with example_context_manager():
    print(1)

with optional(True, example_context_manager()):
    print(2)

with optional(False, example_context_manager()):
    print(3)

Выход:

before
1
after
before
2
after
3

Answer 2

Новые посетители могут быть заинтересованы в contextlib.ExitStack:

      with ExitStack() as stack:
  if condition:
    stack.enter_context(freeze_time(...))
  lines_of_code_1
  lines_of_code_2
  lines_of_code_3

После этого with утверждение, freeze_time имеет значение только при выполнении условия.

Answer 3

Я бы, вероятно, унаследовал от родительского контекстного менеджера и написал бы что-то вроде этого:

class BaseContextManager:
    def __enter__(self):
        print('using Base')
    def __exit__(self, *args, **kwargs):
        print('exiting Base')


class MineContextManager(BaseContextManager):
    def __init__(self, input=None):
        self.input = input

    def __enter__(self):
        if self.input:
            super().__enter__()

    def __exit__(self, *args, **kwargs):
        if self.input:
            super().__exit__()

if __name__ == '__main__':

    with BaseContextManager():
        print('code with base')

    with MineContextManager():
        print('code without base')

    with MineContextManager(input=True):
        print('code again with base')

Это дает:

using Base
code with base
exiting Base
code without base
using Base
code again with base
exiting Base

Answer 4

Просто используйте

(freeze_time if test_dt else (lambda func: contextmanager(func))(lambda dt: (yield)))(test_dt)

Пример:

from contextlib import contextmanager

test_dt = None

@contextmanager
def freeze_time(test_dt):
    print("frozen")
    yield
    print("unfrozen")

with (freeze_time if test_dt else (lambda func: contextmanager(func))(lambda dt: (yield)))(test_dt):
    print("The cold impairs your judgment.")