135

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

norman
  • 5,128
  • 13
  • 44
  • 75

10 Answers10

140

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.


Python 3.7 and above

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 approach is obviously not as flexible as ExitStack, because this is just a binary choice, whereas ExitStack allows you to add as many exiting things as you want, with complicated logic and so on. But this certainly answers the OP's simple requirements.

Mike
  • 19,114
  • 12
  • 59
  • 91
  • 28
    +1, this should be the selected answer. As pointed [here](https://docs.python.org/3/library/contextlib.html?highlight=contextmanager#simplifying-support-for-single-optional-context-managers) it is meant to deal with this kind of problem. Also, it can be used as a nifty one-liner: `with get_stuff() if needs_with() else ExitStack() as gs`. – farsil Feb 20 '17 at 14:14
  • 2
    IMO better to use the purpose-built `nullcontext` for this use-case: `with get_stuff() if needs_with() else nullcontext() as gs`. `ExitStack` is for more complex programmatic control of entering and exiting contexts, i.e. in docs, "this is a relatively low level API", which is overkill here. – Kache Aug 02 '22 at 06:01
  • @Kache That's a nice, simple answer; I'll add it as a possibility. I'll also point out that `nullcontext` was only added in python 3.7, which came a couple years after this answer. :) – Mike Aug 02 '22 at 10:50
75

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

jamesdlin
  • 81,374
  • 13
  • 159
  • 204
  • 6
    This context manager should be in the standard python library, imho. Thanks for this. – jjmontes Jan 21 '16 at 18:37
  • How about using the names *should_...* and *dont*. Then statements like this would read `with get_stuff() if should_get_stuff() else dont() as gs:` ? – Rian Rizvi Mar 25 '16 at 15:12
  • 1
    @RiazRizvi I wouldn't have personally named it that way; I was using the names from the question. – jamesdlin Mar 25 '16 at 18:01
  • 2
    @jjmontes [contextlib.ExitStack](https://docs.python.org/3/library/contextlib.html#simplifying-support-for-single-optional-context-managers) (new python 3.3) can be used as a dummy context manager. – Bryce Guinta May 14 '17 at 02:48
74

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

Daniel Porteous
  • 5,536
  • 3
  • 25
  • 44
  • 6
    This raises the question though, is it always safe to call get_stuff() *before* entering the with statement? Is `with open(file) as fh` equivalent to `f = open(file)` followed by `with f as fh`? – Caleb Stanford Jun 22 '20 at 12:29
  • 2
    It depends on what your context manager does. Most context managers should not do stuff in their `__init__` method and only do things on their `__enter__` (or `__aenter__`) method, which is called when used in the `with` statement. So the answer unfortunately is "it depends". If you're worried about it, you could instead assign the functions to `cm` without calling them (with `functools.partial` if necessary) and then do `with cm() as gs`. – Daniel Porteous Oct 14 '20 at 16:30
  • 1
    One-liner would be: `with get_stuff() if needs_with() else nullcontext() as gs` – nirvana-msu Apr 06 '21 at 00:21
11

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
Anentropic
  • 32,188
  • 12
  • 99
  • 147
  • Does it support an `as ...` clause at the end of the `with` statement? – Craig McQueen Dec 09 '16 at 00:40
  • 1
    looking at the source... yes it does. `with conditional(needs_with(), get_stuff()) as stuff:` will give you a reference to the `get_stuff()` context manager (if and only if the condition is met, otherwise you get `None`) – Anentropic Dec 09 '16 at 11:27
  • I have found your answer is incomplete: https://stackoverflow.com/questions/27803059/conditional-with-statement-in-python/58578167#58578167 – Andry Oct 27 '19 at 09:37
6
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')
Lucas Vazquez
  • 1,456
  • 16
  • 20
4

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
tdelaney
  • 73,364
  • 6
  • 83
  • 116
  • 1
    As I mention in my answer, python 3.3 has added `contextlib.ExitStack`, which appears to do very much what your `ContextGroup` does. I will say I'm a little surprised that it hasn't been backported, but if you're willing to require python >=3.3, that might be a nice robust alternative for you. – Mike Jan 14 '16 at 20:02
  • 1
    `contextlib2` is a pypi package which has backported `ExitStack` to python 2 – anthony sottile Jul 03 '17 at 16:27
3

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.

skeller88
  • 4,276
  • 1
  • 32
  • 34
0

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

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.

Joooeey
  • 3,394
  • 1
  • 35
  • 49
-2

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__'
Andry
  • 2,273
  • 29
  • 28
  • 1
    you're still only passing two arguments to `conditional`, this is not adding anything to my original answer... it looks like your issue was just in correctly preparing the two arguments needed – Anentropic Oct 28 '19 at 11:43
  • 1
    @Anentropic My point is that it is not enough to just put `conditional` to work, additionally needs worksarounds to make it work which not that simple like in your answer and makes the whole thing less applicaible. – Andry Oct 29 '19 at 00:59