Conditional or optional context managers in with statement

Question

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

Could you give a less abstract example? Yes, you could alter the context manager to not do whatever it's doing in enter and exit `if input is None:`. Have you tested what you've written? — jonrsharpe, Dec 20 '16 at 21:57
maybe make a `blank_context_manager` object that does nothing and override `context_manager.__new__` to return the blank one if there is no input, would reduce the number of conditionals to 1. — Tadhg McDonald-Jensen, Dec 20 '16 at 21:59
Possible duplicate of [Conditional with statement in Python](http://stackoverflow.com/questions/27803059/conditional-with-statement-in-python) — Tadhg McDonald-Jensen, Dec 20 '16 at 22:02
@jonrsharpe What part of the code would you like to be less abstract? I'd rather not alter or touch the original context manager. I'll go ahead and replace the code with the actual context manager I'm using, but I feel like the concept itself could apply to more general situations. — Jordan Reiter, Dec 20 '16 at 22:02
I agree that this appears to be a duplicate question but I think Alex Hall's answer is better then the ones on the other page… I'll mark it as the answer. — Jordan Reiter, Dec 20 '16 at 22:11
Does this answer your question? [Is it possible to have an optional with/as statement in python?](https://stackoverflow.com/questions/12168208/is-it-possible-to-have-an-optional-with-as-statement-in-python) — Lucas Cimon, Jul 22 '20 at 10:18

score 15 · Accepted Answer · answered Dec 20 '16 at 22:03

15

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

answered Dec 20 '16 at 22:03

Alex Hall

34,833
5
57
89

I agree with others that this is a duplicate question, but I like your answer much better than the ones already there! Maybe post it on the other page as an answer? http://stackoverflow.com/questions/27803059/conditional-with-statement-in-python – Jordan Reiter Dec 20 '16 at 22:09
If using Python 3.7+, use `nullcontext` in stdlib instead, which is purpose-built for this use-case: https://docs.python.org/3/library/contextlib.html#contextlib.nullcontext – Kache Aug 02 '22 at 06:04
@Kache thanks, I'd never heard of `nullcontext` before! – Alex Hall Aug 02 '22 at 10:28

score 6 · Answer 2 · answered May 03 '21 at 04:38

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.

score 1 · Answer 3 · answered Dec 20 '16 at 22:07

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

cowlinator · Answer 4 · 2020-03-04T01:16:01.270

0

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

edited Mar 04 '20 at 01:16

answered Mar 04 '20 at 01:09

cowlinator

7,195
6
41
61

Conditional or optional context managers in with statement

4 Answers4

Linked