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)  
Was it helpful?

Solution

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.

OTHER TIPS

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)

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

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.

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})
Licensed under: CC-BY-SA with attribution
Not affiliated with StackOverflow
scroll top