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

Figured it out using the contextmanager decorator:

import os
from contextlib import ExitStack, contextmanager

@contextmanager
def profile_if_enabled():
    with ExitStack() as stack:
        if os.environ.get("SHOULD_PROFILE"):
            stack.enter_context(profile())
        yield