Search code examples
pythondecoratornewrelic

Python, conditional use of a decorator (package not installed in all env)


I am using New Relic to monitor my Python (2.7) app, but only in my Production environment. I want to use their function_trace decorator.

In wsgi.py

try:
    import newrelic.agent
    newrelic.agent.initialize('/path/')
except Exception as e:
    pass

In views.py

if settings.ENV == "test":
    pass
else:
    import newrelic

@newrelic.agent.function_trace()
def my_function():
    ....

This works great in Production, but fails in Test, of course. I can't put if-else directly around a decorator so I was thinking that I could use a conditional decorator:

def newrelic_conditional_decorator(function)
    # What magic goes in here...
    if settings.ENV == "test":
        just return function
    else:
        use @newrelic.agent.function_trace() decorator

@newrelic_conditional_decorator
def my_function():
    ....

I'm always a bit hazy on decorators so am hoping for some help here! (Or another way to handle not having the same package in Test as in Production.)


Solution

  • A decorator takes in a function and returns a new function from it. So if you want a conditional decorator, all you need is to return the initial function when you do not want the decorator to be applied.

    Your specific case

    def newrelic_conditional_decorator(function):
    
        if settings.ENV == "test":
            # We do not apply the decorator to the function
            return function
        else:
            # We apply the decorator and return the new function
            return newrelic.agent.function_trace()(function)
    
    @newrelic_conditional_decorator
    def my_function():
        ...
    

    A bit of generality

    You could abstract this to make a more general conditional decorator which applies a given generator provided some condition function returns True.

    def conditional_decorator(decorator, condition, *args):
    
        def wrapper(function):
            if condition(*args):
                return decorator(function)
            else:
                return function
    
        return wrapper
    
    @conditional_decorator(
        newrelic.agent.function_trace(),
        lambda settings: settings.ENV != "test",
        settings)
    def my_function():
        ...