how to do a conditional decorator in python

Question

Is it possible to decorator a function conditionally. For example, I want to decorate the function foo() with a timer function (timeit) only doing_performance_analysis is True (see the psuedo-code below).

if doing_performance_analysis:
  @timeit
  def foo():
    """
    do something, timeit function will return the time it takes
    """
    time.sleep(2)
else:
  def foo():
    time.sleep(2)  

Solution 1

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.

Solution 2

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)

Solution 3

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():
   ...

Solution 4

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()

Solution 5

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.

View more solutions

Comments

[cfpete]

Thanks! The comment section does not format, so I added the sample code to your original reply. Can you pls explain why the timing function is not called?

[Martijn Pieters]

The decorator is applied at import time, so instance variables are not being consulted at that time. You'd have to write a different decorator for that, one that inspects self when called. Out of scope for this Q and comment format. :-)

[PythonEnthusiast]

How would I use it if I want to use this on class based methods, ie Django Class based views. Over there we have to use method_decorator. How to make this code compatible with that?

[Martijn Pieters]

@PythonEnthusiast: By just passing in the result of the call: @method_decorator(conditional_decorator(timeit, doing_performance_analysis)).