Conditional Or Optional Context Managers In With Statement

Suppose I have some kind of context manager (from a third-party library) that I am using like so: with freeze_time(test_dt): lines_of_code_1 lines_of_code_2 lines_of_co

Solution 1:

Here's an easy way to wrap around an existing context manager without even using any classes:

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)

Output:

before
1
after
before
2
after
3

Solution 2:

Newer visitors may be interested in 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

After this with statement, freeze_time is only relevant on the condition being true.

Solution 3:

I'd probably inherit from the parent context manager and write something like this:

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')

This gives:

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

Solution 4:

Just use

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

Example:

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.")