How to do a conditional decorator in python?

Question:

Is it possible to decorate a function conditionally? For example, I want to decorate the function foo() with a timer function (timeit), but only when doing_performance_analysis condition is True, like this:

    if doing_performance_analysis:
      @timeit
      def foo():
        """
        Do something, e.g. sleep, and let timeit 
        return the time it takes
        """
        time.sleep(2)
    else:
      def foo():
        time.sleep(2)
Asked By: cfpete

||
Source

Answers:

How about:

def foo():
   ...

if doing_performance_analysis:
   foo = timeit(foo)

I imagine you could even wrap this into a decorator that would take a boolean flag and another decorator, and would only apply the latter if the flag is set to True:

def cond_decorator(flag, dec):
   def decorate(fn):
      return dec(fn) if flag else fn
   return decorate

@cond_decorator(doing_performance_analysis, timeit)
def foo():
   ...
Answered By: NPE

A decorator is simply a function applied to another function. You can apply it manually:

def foo():
   # whatever
   time.sleep(2)

if doing_performance_analysis:
    foo = timeit(foo)
Answered By: Blckknght

Decorators are simply callables that return a replacement, optionally the same function, a wrapper, or something completely different. As such, you could create a conditional decorator:

def conditional_decorator(dec, condition):
    def decorator(func):
        if not condition:
            # Return the function unchanged, not decorated.
            return func
        return dec(func)
    return decorator

Now you can use it like this:

@conditional_decorator(timeit, doing_performance_analysis)
def foo():
    time.sleep(2)

The decorator could also be a class:

class conditional_decorator(object):
    def __init__(self, dec, condition):
        self.decorator = dec
        self.condition = condition

    def __call__(self, func):
        if not self.condition:
            # Return the function unchanged, not decorated.
            return func
        return self.decorator(func)

Here the __call__ method plays the same role as the returned decorator() nested function in the first example, and the closed-over dec and condition parameters here are stored as arguments on the instance until the decorator is applied.

Answered By: Martijn Pieters

Blckknght’s answer is great if you want to do the check every time you call the function, but if you have a setting that you can read once and never changes you may not want to check the setting every time the decorated function is called. In some of our high performance daemons at work I have written a decorator that checks a setting file once when the python file is first loaded and decides if it should wrap it or not.

Here is a sample

def timed(f):
    def wrapper(*args, **kwargs):
        start = datetime.datetime.utcnow()
        return_value = f(*args, **kwargs)
        end = datetime.datetime.utcnow()
        duration = end - start

        log_function_call(module=f.__module__, function=f.__name__, start=__start__, end=__end__, duration=duration.total_seconds())
    if config.get('RUN_TIMED_FUNCTIONS'):
        return wrapper
    return f

Assuming that log_function_call logs your call to a database, logfile, or whatever and that config.get(‘RUN_TIMED_FUNCTIONS’) checks your global configuration, then adding the @timed decorator to a function will check once on load to see if you are timing on this server, environment, etc. and if not then it won’t change the execution of the function on production or the other environments where you care about performance.

Answered By: nobled
use_decorator = False

class myDecorator(object):
    def __init__(self, f):
            self.f = f

    def __call__(self):
            print "Decorated running..."
            print "Entering", self.f.__name__
            self.f()
            print "Exited", self.f.__name__


def null(a):
    return a


if use_decorator == False :
    myDecorator = null


@myDecorator
def CoreFunction():
    print "Core Function running"

CoreFunction()
Answered By: Lawrence Chernin

Here is what worked for me:

def timeit(method):
    def timed(*args, **kw):
        if 'usetimer' not in kw:
            return method(*args, **kw)
        elif ('usetimer' in kw and kw.get('usetimer') is None):
            return method(*args, **kw)
        else:
            import time
            ts = time.time()
            result = method(*args, **kw)
            te = time.time()
            if 'log_time' in kw:
                name = kw.get('log_name', method.__name__.upper())
                kw['log_time'][name] = int((te - ts) * 1000)
            else:
                print '%r took %2.2f ms' % 
                      (method.__name__, (te - ts) * 1000)
            return result
    return timed

def some_func(arg1, **kwargs):
    #do something here

some_func(param1, **{'usetimer': args.usetimer})
Answered By: diman82

A bit late to the party but this post was very helpful for me and I wanted to add something I found useful. I wanted to do a conditional decorator where the decorator took some arguments. You can do this very cleanly with partial.

from functools import partial

def conditional_decorator(dec, condition, **kwargs):
    def decorator(func):
        if not condition:
            # Return the function unchanged, not decorated.
            return func  # train
        return partial(dec, **kwargs)(func)
    return decorator

This lets you partially apply the arguments that your decorator takes and still make it conditionally apply.

Take the following example where we have a decorator that adds n to the result of some function:

def add_n(func, n=1):
    def wrapper_add_n(*args, **kwargs):
        return func(*args, **kwargs) + n
    return wrapper_add_n

We can use the modified conditional decorator to add in the n parameters.


@conditional_decorator(add_n, True, n=3)
def multiply(a, b):
    return a * b

print(multiply(1, 2))

Which prints 5.

Answered By: nbertagnolli