Use the same function as context manager and decorator

I recently learned that context managers created with @contextmanager can be used either as a context manager or a decorator:

from contextlib import contextmanager

@contextmanager
def test_context():
    print('Entering')
    yield
    print('Leaving')

with test_context():
    print('Inside')

@test_context()
def test_decorated():
    print('Decorated')

test_decorated()

We will yield:

Entering
Inside
Leaving

Entering
Decorated
Leaving

In both cases, we can pass argument to our context manager/decorator like this:

from contextlib import contextmanager

@contextmanager
def test_context(value):
    print('Entering:', value)
    yield
    print('Leaving')

with test_context('Hello there!'):
    print('Inside')

@test_context('Hello there!')
def test_decotared():
    print('Decorated')

But what if we need to access the arguments of the function inside the decorator? When used as a context manager, we have to pass them, but we cannot access them directly when we use it as a decorator. To do this, we need to take inspiration from contextlib.ContextDecorator (what contextlib.contextmanager is using) to to some work ourselves:

from functools import wraps

class test_context:
    def __init__(self, option_value=None, function_value=None):
        self._option_value = option_value
        self._function_value = function_value

    def __enter__(self):
        # Called when entering the ctx manager.
        print('Entering:', self._option_value, self._function_value)
        return self

    def __exit__(self, *exc):
        # Called when leaving the ctx manager.
        print('Leaving')

    def __call__(self, func):
        # Called when used as a decorator.
        @wraps(func)
        def wrapper(function_value):
            self._function_value = function_value
            with self._recreate_cm():
                return func(function_value)

        return wrapper

    def _recreate_cm(self):
        # Taken as is from contextlib.ContextDecorator.
        return self

with test_context('Option value', 'Function value'):
    print('Inside')

@test_context('Option value')
def test_decorated(value):
    print('Decorated', value)

test_decorated('Function value')

This will yield:

Entering: Option value Function value
Inside
Leaving

Entering: Option value Function value
Decorated Function value
Leaving