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

Solution

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