Conditional with statement in Python

Question

Is there a way to begin a block of code with a with statement, but conditionally?

Something like:

if needs_with():
    with get_stuff() as gs:

# do nearly the same large block of stuff,
# involving gs or not, depending on needs_with()

To clarify, one scenario would have a block encased in the with statement, while another possibility would be the same block, but not encased (i.e., as if it wasn't indented)

Initial experiments of course give indentation errors..

Answer 1

You can use contextlib.nested to put 0 or more context managers into a single with statement.

>>> import contextlib
>>> managers = []
>>> test_me = True
>>> if test_me:
...     managers.append(open('x.txt','w'))
... 
>>> with contextlib.nested(*managers):                                                       
...  pass                                                    
...                                                             
>>> # see if it closed
... managers[0].write('hello')                                                                                                                              
Traceback (most recent call last):                              
  File "<stdin>", line 2, in <module>                                   
ValueError: I/O operation on closed file

This solution has its quirks and I just noticed that as of 2.7 its been deprecated. I wrote my own context manager to handle juggling multiple context managers. Its worked for me so far, but I haven't really considered edge conditons

class ContextGroup(object):
    """A group of context managers that all exit when the group exits."""

    def __init__(self):
        """Create a context group"""
        self._exits = []

    def add(self, ctx_obj, name=None):
        """Open a context manager on ctx_obj and add to this group. If
        name, the context manager will be available as self.name. name
        will still reference the context object after this context
        closes.
        """
        if name and hasattr(self, name):
            raise AttributeError("ContextGroup already has context %s" % name)
        self._exits.append(ctx_obj.__exit__)
        var = ctx_obj.__enter__()
        if name:
            self.__dict__[name] = var

    def exit_early(self, name):
        """Call __exit__ on named context manager and remove from group"""
        ctx_obj = getattr(self, name)
        delattr(self, name)
        del self._exits[self._exits.index(ctx_obj)]
        ctx_obj.__exit__(None, None, None)

    def __enter__(self):
        return self

    def __exit__(self, _type, value, tb):
        inner_exeptions = []
        for _exit in self._exits:
            try:
                _exit(_type, value, tb )
            except Exception, e:
                inner_exceptions.append(e)
        if inner_exceptions:
            r = RuntimeError("Errors while exiting context: %s" 
                % (','.join(str(e)) for e in inner_exceptions))

    def __setattr__(self, name, val):
        if hasattr(val, '__exit__'):
            self.add(val, name)
        else:
            self.__dict__[name] = val

Answer 2

If you want to avoid duplicating code and are using a version of Python prior to 3.7 (when contextlib.nullcontext was introduced) or even 3.3 (when contextlib.ExitStack was introduced), you could do something like:

class dummy_context_mgr():
    def __enter__(self):
        return None
    def __exit__(self, exc_type, exc_value, traceback):
        return False

or:

import contextlib

@contextlib.contextmanager
def dummy_context_mgr():
    yield None

and then use it as:

with get_stuff() if needs_with() else dummy_context_mgr() as gs:
   # do stuff involving gs or not

You alternatively could make get_stuff() return different things based on needs_with().

See Mike or Daniel's answers for what you can do in later versions

Answer 3

Python 3.3 introduced contextlib.ExitStack for just this kind of situation. It gives you a "stack", to which you add context managers as necessary. In your case, you would do this:

from contextlib import ExitStack

with ExitStack() as stack:
    if needs_with():
        gs = stack.enter_context(get_stuff())

    # do nearly the same large block of stuff,
    # involving gs or not, depending on needs_with()

Anything that is entered to stack is automatically exited at the end of the with statement as usual. (If nothing is entered, that's not a problem.) In this example, whatever is returned by get_stuff() is exited automatically.

If you have to use an earlier version of python, you might be able to use the contextlib2 module, although this is not standard. It backports this and other features to earlier versions of python. You could even do a conditional import, if you like this approach.

Answer 4

It was hard to find @farsil's nifty Python 3.3 one-liner, so here it is in its own answer:

with ExitStack() if not needs_with() else get_stuff() as gs:
     # do stuff

Note that ExitStack should come first, otherwise get_stuff() will be evaluated.

Answer 5

A third-party option to achieve exactly this:
https://pypi.python.org/pypi/conditional

from conditional import conditional

with conditional(needs_with(), get_stuff()):
    # do stuff

Answer 6

As of Python 3.7 you can use contextlib.nullcontext:

from contextlib import nullcontext

if needs_with():
    cm = get_stuff()
else:
    cm = nullcontext()

with cm as gs:
    # Do stuff

contextlib.nullcontext is pretty much just a no-op context manager. You can pass it an argument that it will yield, if you depend on something existing after the as:

>>> with nullcontext(5) as value:
...     print(value)
...
5

Otherwise it'll just return None:

>>> with nullcontext() as value:
...     print(value)
...
None

It's super neat, check out the docs for it here: https://docs.python.org/3/library/contextlib.html#contextlib.nullcontext