Search code examples
pythonpython-3.xwith-statementcontextmanager

Using different context managers depending on condition


Is it possible to execute a single block using different context managers depending on some condition?

Example:

if some_condition:
    with ContextManager(**args) as contex:
        ... # some block
else:
    with OtherContextManager(**other_args) as contex:
        ... # the same block

One way would be to wrap ... into a function, but this may not be too convenient in my case. Is there another possibility?


Solution

  • You can store the constructed object in a variable, like:

    if some_condition:
        cm = ContextManager(**args)
    else:
        cm = OtherContextManager(**other_args)
    
    with cm as contex:
            ... # some block
    

    The above can easily be extended to three possible context managers, etc. You can also decide for example to first "patch" the context manager before "entering" the context.

    Although it is common to see a pattern like with foo() as bar:, in fact Python simply evaluates the foo(), obtains that element, and calls .__enter__() on the object. The result of that method is stored in the bar.

    So there is nothing "special" about the foo() call, you can use any kind of object on the left side. You can thus for example encapsulate the if-else logic in a separate function, and return the context manager, and then use the variable, or pass context managers as parameters. As long as you use it in a with statement, Python will call .__enter__(..) and .__exit__(..) behind the curtains.