How do I make function decorators and chain them together?

Viewed 616299

How do I make two decorators in Python that would do the following?

@make_bold
@make_italic
def say():
   return "Hello"

Calling say() should return:

"<b><i>Hello</i></b>"
18 Answers

Check out the documentation to see how decorators work. Here is what you asked for:

from functools import wraps

def makebold(fn):
    @wraps(fn)
    def wrapper(*args, **kwargs):
        return "<b>" + fn(*args, **kwargs) + "</b>"
    return wrapper

def makeitalic(fn):
    @wraps(fn)
    def wrapper(*args, **kwargs):
        return "<i>" + fn(*args, **kwargs) + "</i>"
    return wrapper

@makebold
@makeitalic
def hello():
    return "hello world"

@makebold
@makeitalic
def log(s):
    return s

print hello()        # returns "<b><i>hello world</i></b>"
print hello.__name__ # with functools.wraps() this returns "hello"
print log('hello')   # returns "<b><i>hello</i></b>"

Alternatively, you could write a factory function which return a decorator which wraps the return value of the decorated function in a tag passed to the factory function. For example:

from functools import wraps

def wrap_in_tag(tag):
    def factory(func):
        @wraps(func)
        def decorator():
            return '<%(tag)s>%(rv)s</%(tag)s>' % (
                {'tag': tag, 'rv': func()})
        return decorator
    return factory

This enables you to write:

@wrap_in_tag('b')
@wrap_in_tag('i')
def say():
    return 'hello'

or

makebold = wrap_in_tag('b')
makeitalic = wrap_in_tag('i')

@makebold
@makeitalic
def say():
    return 'hello'

Personally I would have written the decorator somewhat differently:

from functools import wraps

def wrap_in_tag(tag):
    def factory(func):
        @wraps(func)
        def decorator(val):
            return func('<%(tag)s>%(val)s</%(tag)s>' %
                        {'tag': tag, 'val': val})
        return decorator
    return factory

which would yield:

@wrap_in_tag('b')
@wrap_in_tag('i')
def say(val):
    return val
say('hello')

Don't forget the construction for which decorator syntax is a shorthand:

say = wrap_in_tag('b')(wrap_in_tag('i')(say)))

Decorators are just syntactical sugar.

This

@decorator
def func():
    ...

expands to

def func():
    ...
func = decorator(func)

Python decorators add extra functionality to another function

An italics decorator could be like

def makeitalic(fn):
    def newFunc():
        return "<i>" + fn() + "</i>"
    return newFunc

Note that a function is defined inside a function. What it basically does is replace a function with the newly defined one. For example, I have this class

class foo:
    def bar(self):
        print "hi"
    def foobar(self):
        print "hi again"

Now say, I want both functions to print "---" after and before they are done. I could add a print "---" before and after each print statement. But because I don't like repeating myself, I will make a decorator

def addDashes(fn): # notice it takes a function as an argument
    def newFunction(self): # define a new function
        print "---"
        fn(self) # call the original function
        print "---"
    return newFunction
    # Return the newly defined function - it will "replace" the original

So now I can change my class to

class foo:
    @addDashes
    def bar(self):
        print "hi"

    @addDashes
    def foobar(self):
        print "hi again"

For more on decorators, check http://www.ibm.com/developerworks/linux/library/l-cpdecor.html

This answer has long been answered, but I thought I would share my Decorator class which makes writing new decorators easy and compact.

from abc import ABCMeta, abstractclassmethod

class Decorator(metaclass=ABCMeta):
    """ Acts as a base class for all decorators """

    def __init__(self):
        self.method = None

    def __call__(self, method):
        self.method = method
        return self.call

    @abstractclassmethod
    def call(self, *args, **kwargs):
        return self.method(*args, **kwargs)

For one I think this makes the behavior of decorators very clear, but it also makes it easy to define new decorators very concisely. For the example listed above, you could then solve it as:

class MakeBold(Decorator):
    def call():
        return "<b>" + self.method() + "</b>"

class MakeItalic(Decorator):
    def call():
        return "<i>" + self.method() + "</i>"

@MakeBold()
@MakeItalic()
def say():
   return "Hello"

You could also use it to do more complex tasks, like for instance a decorator which automatically makes the function get applied recursively to all arguments in an iterator:

class ApplyRecursive(Decorator):
    def __init__(self, *types):
        super().__init__()
        if not len(types):
            types = (dict, list, tuple, set)
        self._types = types

    def call(self, arg):
        if dict in self._types and isinstance(arg, dict):
            return {key: self.call(value) for key, value in arg.items()}

        if set in self._types and isinstance(arg, set):
            return set(self.call(value) for value in arg)

        if tuple in self._types and isinstance(arg, tuple):
            return tuple(self.call(value) for value in arg)

        if list in self._types and isinstance(arg, list):
            return list(self.call(value) for value in arg)

        return self.method(arg)


@ApplyRecursive(tuple, set, dict)
def double(arg):
    return 2*arg

print(double(1))
print(double({'a': 1, 'b': 2}))
print(double({1, 2, 3}))
print(double((1, 2, 3, 4)))
print(double([1, 2, 3, 4, 5]))

Which prints:

2
{'a': 2, 'b': 4}
{2, 4, 6}
(2, 4, 6, 8)
[1, 2, 3, 4, 5, 1, 2, 3, 4, 5]

Notice that this example didn't include the list type in the instantiation of the decorator, so in the final print statement the method gets applied to the list itself, not the elements of the list.

Paolo Bergantino's answer has the great advantage of only using the stdlib, and works for this simple example where there are no decorator arguments nor decorated function arguments.

However it has 3 major limitations if you want to tackle more general cases:

  • as already noted in several answers, you can not easily modify the code to add optional decorator arguments. For example creating a makestyle(style='bold') decorator is non-trivial.
  • besides, wrappers created with @functools.wraps do not preserve the signature, so if bad arguments are provided they will start executing, and might raise a different kind of error than the usual TypeError.
  • finally, it is quite difficult in wrappers created with @functools.wraps to access an argument based on its name. Indeed the argument can appear in *args, in **kwargs, or may not appear at all (if it is optional).

I wrote decopatch to solve the first issue, and wrote makefun.wraps to solve the other two. Note that makefun leverages the same trick than the famous decorator lib.

This is how you would create a decorator with arguments, returning truly signature-preserving wrappers:

from decopatch import function_decorator, DECORATED
from makefun import wraps

@function_decorator
def makestyle(st='b', fn=DECORATED):
    open_tag = "<%s>" % st
    close_tag = "</%s>" % st

    @wraps(fn)
    def wrapped(*args, **kwargs):
        return open_tag + fn(*args, **kwargs) + close_tag

    return wrapped

decopatch provides you with two other development styles that hide or show the various python concepts, depending on your preferences. The most compact style is the following:

from decopatch import function_decorator, WRAPPED, F_ARGS, F_KWARGS

@function_decorator
def makestyle(st='b', fn=WRAPPED, f_args=F_ARGS, f_kwargs=F_KWARGS):
    open_tag = "<%s>" % st
    close_tag = "</%s>" % st
    return open_tag + fn(*f_args, **f_kwargs) + close_tag

In both cases you can check that the decorator works as expected:

@makestyle
@makestyle('i')
def hello(who):
    return "hello %s" % who

assert hello('world') == '<b><i>hello world</i></b>'    

Please refer to the documentation for details.

I add a case when you need to add custom parameters in decorator, pass it to final function and then work it with.

the very decorators:

def jwt_or_redirect(fn):
  @wraps(fn)
  def decorator(*args, **kwargs):
    ...
    return fn(*args, **kwargs)
  return decorator

def jwt_refresh(fn):
  @wraps(fn)
  def decorator(*args, **kwargs):
    ...
    new_kwargs = {'refreshed_jwt': 'xxxxx-xxxxxx'}
    new_kwargs.update(kwargs)
    return fn(*args, **new_kwargs)
  return decorator

and the final function:

@app.route('/')
@jwt_or_redirect
@jwt_refresh
def home_page(*args, **kwargs):
  return kwargs['refreched_jwt']

Yet another example of nested decorators for plotting an image:

import matplotlib.pylab as plt

def remove_axis(func):
    def inner(img, alpha):
        plt.axis('off')
        func(img, alpha)
    return inner

def plot_gray(func):
    def inner(img, alpha):
        plt.gray()
        func(img, alpha)
    return inner

@remove_axis
@plot_gray
def plot_image(img, alpha):
    plt.imshow(img, alpha=alpha)
    plt.show()

Now, let's show a color image first without axis labels using the nested decorators:

plot_image(plt.imread('lena_color.jpg'), 0.4)

enter image description here

Next, let's show a gray scale image without axis labels using the nested decorators remove_axis and plot_gray (we need to cmap='gray', otherwise the default colormap is viridis, so a grayscale image is by default not displayed in black and white shades, unless explicitly specified)

plot_image(plt.imread('lena_bw.jpg'), 0.8)

enter image description here

The above function call reduces down to the following nested call

remove_axis(plot_gray(plot_image))(img, alpha)
Related