danijar danijar - 2 months ago 20
Python Question

How to reuse a custom @contextlib.contextmanager?

I'm trying to create a wrapper to make context objects optional. When the condition is true, the thing should behave like the wrapped context object, otherwise it should behave like a no-op context object. This example works for using the wrapped object once, but fails it it is reused.

Example:

import contextlib
from threading import Lock

@contextlib.contextmanager
def conditional_context(context, condition):
if condition and context:
with context:
yield
else:
yield

use_locking = False
lock = conditional_context(Lock(), use_locking)
with lock:
print('Foo')
with lock:
print('Bar')


Output:

Foo
Traceback (most recent call last):
File "example.py", line 16, in <module>
with lock:
File "/usr/lib/python3.5/contextlib.py", line 61, in __enter__
raise RuntimeError("generator didn't yield") from None
RuntimeError: generator didn't yield

Answer

You can't do that with contextlib.contextmanager. As mentioned in passing in the docs, context managers created by contextmanager are one-shot.

You will have to write your own class with __enter__ and __exit__ methods if you want the same object to be reusable in multiple with statements:

from threading import Lock


class ConditionalContext:

    def __init__(self, context, condition):
        self._context = context
        self._condition = condition

    def __enter__(self, *args, **kwargs):
        if self._condition:
            self._context.__enter__(*args, **kwargs)

    def __exit__(self, *args, **kwargs):
        if self._condition:
            self._context.__exit__(*args, **kwargs)


use_locking = False
lock = ConditionalContext(Lock(), use_locking)
with lock:
    print('Foo')
with lock:
    print('Bar')
Comments