How to decorate a Python function with a context manager?

user_5104777 0
Viewed 458

Given an existing context manager, I'd like a function decorator that causes execution to occur inside a with-block.

This should make the following two code blocks equivalent:

@decorate(contextmanager)
async def f():
   ...

await f()

async def f():
    ...

async with contextmanager:
    await f()

This would be handy for wrapping functions in e.g. asyncio.Semaphore(n) contexts.

Does a convenience utility like decorate already exist somewhere in the standard library?

contextmanager python-decorators python
1 Answers

There is ContextDecorator, you can create context managers that inherit from it and use it like this

from functools import wraps
from contextlib import ContextDecorator

class MyContext(ContextDecorator):
    def __enter__(self):
        print('in enter')
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        print('in exit')

some_context = MyContext()

@some_context
def foo(x, y):
    print(x + y)

foo(5, 3)

However, I'm not sure you can use it for existing things you didn't write, or if it works for async stuff, in that case you can just write a decorator like this:

from functools import wraps

# Some context manager
class MyContext:
    def __enter__(self):
        print('in enter')
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        print('in exit')

# Decotator that gets an existing context manager and uses it
def context_decorator(context_manager):

    def inner(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            with context_manager:
                return func(*args, **kwargs)

        return wrapper

    return inner

some_context = MyContext()

@context_decorator(some_context)
def foo(x, y):
    print(x + y)

foo(5, 3)

in enter
8
in exit

ps. note that for something like asyncio.Semephore you need to modify the code to use async with and await in the relevant places

user_21262540
Related