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_code_3

But, suppose if there is no value for test_dt, the context manager should not run, but all of the remaining code should run, like so:

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

Assume that lines_of_code here is 2-3 lines of code which are exactly identical, is there a cleaner way of writing this? I’m aware that I could write something like this:

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

But I’m not crazy about this formatting. Also, I don’t want to have to litter this pattern all over my code.

There is one final possibility, but I’m not certain it will work: subclassing the context manager and skipping the __enter__ and __exit__ functions if the test_dt given is empty, like so:

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)

I tested it out with a blank context manager class, and it seemed to behave correctly. However, I’m worried whether a real context manager will behave correctly if I do this (I’m not very familiar with the internals of how it works).

Answer

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