Conditional context manager and decorator in Python

Question

I have the following Python code littered throughout my code base which profiles a code block and sends the results to a monitoring solution if the SHOULD_PROFILE environment variable is set.

from contextlib import ExitStack

with ExitStack() as stack:
    if os.environ.get("SHOULD_PROFILE"):
        stack.enter_context(profile())
    ...

I would like to consolidate this snippet into a single context manager / decorator so it can be used as a context manager AND a decorator:

with profile_if_enabled():
    ...

# AND

@profile_if_enabled()
def func():
    ....

This is what I've come up with but it isn't working.

from contextlib import ContextDecorator, ExitStack

class profile_if_enabled(ContextDecorator):
    def __enter__(self):
        with ExitStack() as stack:
            if os.environ.get("SHOULD_PROFILE"):
                stack.enter_context(profile())
            return self

    def __exit__(self, *exc):
        return False

Any idea what I'm doing wrong?

Answer 1

Your current attempt fails because as soon as you return self , you leave the with ExitStack() as stack: block, and the ExitStack performs cleanup. You need to perform cleanup in __exit__ , not when __enter__ returns.

ContextDecorator requires a reusable, ideally reentrant context manager, which doesn't match your needs well. Even if you get something working, it will most likely break if you try to call a profiled function from another profiled function, or if you try to run two profiled functions in different threads at the same time.

It'd be cleaner to just write a decorator manually:

def profile_if_enabled(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        if os.environ.get('SHOULD_PROFILE'):
            with profile():
                return func(*args, **kwargs)
        else:
            return func(*args, **kwargs)
    return wrapper

If you really want something that works as both a decorator and a context manager, you can check for the presence of an argument:

def profile_if_enabled(func=None):
    if func is None:
        return profile()

    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        if os.environ.get('SHOULD_PROFILE'):
            with profile():
                return func(*args, **kwargs)
        else:
            return func(*args, **kwargs)
    return wrapper

This would then be used as either with profile_if_enabled(): or @profile_if_enabled . Note that unlike with ContextDecorator , the use as a decorator does not have parentheses.

If you don't need to support changing the SHOULD_PROFILE setting during execution, you could alternatively restructure the code to avoid repeating the check:

if os.environ.get('SHOULD_PROFILE'):
    def profile_if_enabled(func=None):
        if func is None:
            return profile()

        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            with profile():
                return func(*args, **kwargs)
        return wrapper
else:
    def profile_if_enabled(func=None):
        if func is None:
            return contextlib.nullcontext()
        return func