CODEWITHSUNDEEP
pythonpython-3.xwith-statementcontextmanager
cheersmate
cheersmate

Reputation: 2656

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?

Upvotes: 6

Views: 1216

Answers (3)

yorodm
yorodm

Reputation: 4461

We can just go crazy and take advantage of the fact that both __enter__ and __exit__ are just methods and that they're called in the original object (not the one returned by __enter__):

class WrapperContext:

 def __init__(self, condition):
     if condition:
         self.real_context = ContextA()
     else:
         self.real_context = ContextB()

 def __enter__(self):
     return self.real_context.__enter__()

 def __exit__(self):
     return self.real_context.__exit__()

and use it like this:

 with WrapperContext(condition) as obj:

Upvotes: 2

willeM_ Van Onsem
willeM_ Van Onsem

Reputation: 476584

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.

Upvotes: 7

Martin Stone
Martin Stone

Reputation: 13007

How about...

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

...?

Upvotes: 2

Related Questions