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

Python 3.7 and above

The ExitStack described below will also work, but Python 3.7 further introduced contextlib.nullcontext (a couple years after this answer was originally posted, and since that time mentioned in several other answers). In the comments, @Kache points out the most elegant usage of this option:

from contextlib import nullcontext

with get_stuff() if needs_with() else nullcontext() as gs:
    # do nearly the same large block of stuff,
    # involving gs or not, depending on needs_with()

Note that if needs_with() is False, then gs will be None inside the context block. If you want gs to be something_else in that case, you just replace nullcontext() with nullcontext(something_else).

This simple approach provides just a binary choice between a nullcontext and some other context. While this answers the OP's simple requirements, a more sophisticated is also possible for more complex scenarios: in Python 3.3 and above the ExitStack allows you to add as many exiting things as you want, with complicated logic and so on.

---

Python 3.3 and above

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 2

Encapsulate the body of the with statement in a function

You can do all kinds of complicated tricks as suggested in the other answers but these are difficult to read and add unnecessary complexity. Your program structure will be easier to follow if you encapsulate the body of your optional with block in its own function. E.g.:

def do_stuff(a, *args, **kwargs):
    ...
    return a

a = 1
gs = "already_present_stuff"
if needs_with():
    with get_stuff() as gs:
        a = do_stuff(a, gs=gs)
else:
    a = do_stuff(a, gs=gs)

Which arguments and keyword argument you pass to do_stuff and what you return is completely up to you. You should pass in everything that is accessed within do_stuff and return every name that you change and every new name you create.

Doing things this way is a great way to avoid spaghetti code and hence complaints from pylint (too many lines per function or method) or human code reviewers if that's applicable.

Yes, there are cases where contextlib is required but most of the time just using a function is the way to go.

Answer 3

import contextlib

my_context = None # your context
my_condition = False # your condition

# Option 1 (Recommended)
with my_context if my_condition else contextlib.nullcontext():
    print('hello 1')

# Option 2
with my_context if my_condition else contextlib.ExitStack():
    print('hello 2')

Answer 4

I have found that the @Anentropic answer is incomplete.

from conditional import conditional

a = 1 # can be None

if not a is None:
  b = 1

class WithNone:
  def __enter__(self):
    return self
  def __exit__(self, type, value, tb):
    pass

def foo(x):
  print(x)
  return WithNone()

with conditional(not a is None, foo(b) if not a is None else None):
  print(123)

The complete conditional usage required 3 conditions instead of 1 because of:

1. NameError: name 'b' is not defined in case if not defined a
2. the function foo still must return enterable object, otherwise: AttributeError: 'NoneType' object has no attribute '__enter__'

Answer 5

So I made this code; It is invoked like so:

with c_with(needs_with(), lambda: get_stuff()) as gs:
    ##DOESN't call get_stuff() unless needs_with is called.
    # do nearly the same large block of stuff,
    # involving gs or not, depending on needs_with()

Properties:

1. it does not call get_stuff() unless condition is true
2. if condition is false, it provides a dummy contextmanager. (could probably be replaced with contextlib.nullcontext for python >= 3.7)
3. Optionally you can send in an alternative contextmanager in case the condition is false:
   with c_with(needs_with(), lambda: get_stuff(), lambda: dont_get_stuff()) as gs:

Hope this will help someone!

-- Here is the code:

def call_if_lambda(f):
    """
    Calls f if f is a lambda function.
    From https://stackoverflow.com/a/3655857/997253
    """
    LMBD = lambda:0
    islambda=isinstance(f, type(LMBD)) and f.__name__ == LMBD.__name__
    return f() if islambda else f
import types
class _DummyClass(object):
    """
    A class that doesn't do anything when methods are called, items are set and get etc.
    I suspect this does not cover _all_ cases, but many.
    """
    def _returnself(self, *args, **kwargs):
        return self
    __getattr__=__enter__=__exit__=__call__=__getitem__=_returnself
    def __str__(self):
        return ""
    __repr__=__str__
    def __setitem__(*args,**kwargs):
        pass
    def __setattr__(*args,**kwargs):
        pass

class c_with(object):
    """
    Wrap another context manager and enter it only if condition is true.
    Parameters
    ----------
    condition:  bool
        Condition to enter contextmanager or possibly else_contextmanager
    contextmanager: contextmanager, lambda or None
        Contextmanager for entering if condition is true. A lambda function
        can be given, which will not be called unless entering the contextmanager.
    else_contextmanager: contextmanager, lambda or None
        Contextmanager for entering if condition is true. A lambda function
        can be given, which will not be called unless entering the contextmanager.
        If None is given, then a dummy contextmanager is returned.
    """
    def __init__(self, condition, contextmanager, else_contextmanager=None):
        self.condition = condition
        self.contextmanager = contextmanager
        self.else_contextmanager = _DummyClass() if else_contextmanager is None else else_contextmanager
    def __enter__(self):
        if self.condition:
            self.contextmanager=call_if_lambda(self.contextmanager)
            return self.contextmanager.__enter__()
        elif self.else_contextmanager is not None:
            self.else_contextmanager=call_if_lambda(self.else_contextmanager)
            return self.else_contextmanager.__enter__()
    def __exit__(self, *args):
        if self.condition:
            return self.contextmanager.__exit__(*args)
        elif self.else_contextmanager is not None:
            self.else_contextmanager.__exit__(*args)

#### EXAMPLE BELOW ####

from contextlib import contextmanager

def needs_with():
    return False

@contextmanager
def get_stuff():
    yield {"hello":"world"}

with c_with(needs_with(), lambda: get_stuff()) as gs:
    ## DOESN't call get_stuff() unless needs_with() returns True.
    # do nearly the same large block of stuff,
    # involving gs or not, depending on needs_with()
    print("Hello",gs['hello'])

Answer 6

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's answer or Daniel's answer for what you can do in later versions.)

Answer 7

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 8

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

Answer 9

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 10

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