Search code examples
pythonconditional-statementsdecoratorpython-decorators

How to do a conditional decorator in python?


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)  

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.