How to make a decorator

Transcript

Let's make a decorator.

We're going to make a function decorator: that is a decorator meant for decorating a function (not for decorating a class).

What the decorator syntax does

We have a decorator function log_me, which is a function decorator (it's used for decorating functions):

def log_me(func):
    def wrapper(*args, **kwargs):
        print("Calling with", args, kwargs)
        return_value = func(*args, **kwargs)
        print("Returning", return_value)
        return return_value
    return wrapper

This decorator can be used like this:

>>> @log_me
... def greet(name):
...     print("Hello", name)
...

That @log_me syntax is a way to define function (greet) that is decorated by this log_me decorator.

This @ syntax is equivalent to defining the greet function first (undecorated):

>>> def greet(name):
...     print("Hello", name)
...

And then taking the variable that points to that function (the variable is just the function name) and pointing it instead to the return value we get when we call our decorator (log_me) with the original function object (the greet function).

>>> greet = log_me(greet)

Yes, this does look a little weird.

A decorator is a function that accepts a function and returns a function.

So we're calling log_me with our greet function object and we're pointing the greet variable to the new function we get back.

So this:

>>> def greet(name):
...     print("Hello", name)
...
>>> greet = log_me(greet)

Is exactly equivalent to this decorator syntax:

>>> @log_me
... def greet(name):
...     print("Hello", name)
...

What happens when we call this decorated function?

We're passing a function to the log_me decorator and then replacing our greet variable with whatever function log_me returns to us.

What do you think we'll see when we call our greet function?

>>> greet("Trey")

Remember, this is the code for our log_me decorator:

def log_me(func):
    def wrapper(*args, **kwargs):
        print("Calling with", args, kwargs)
        return_value = func(*args, **kwargs)
        print("Returning", return_value)
        return return_value
    return wrapper

And this is our decorated greet function:

>>> @log_me
... def greet(name):
...     print("Hello", name)
...

So what's your guess that we'll see when we pass the string Trey to our greet function?

>>> greet("Trey")

When we call greet with the string Trey we see this:

>>> greet("Trey")
Calling with ('Trey',) {}
Hello Trey
Returning None

First Calling with ('Trey',) {} is printed out (that's a tuple and a dictionary of our function arguments), followed by Hello Trey, and then Returning None is printed out.

Why did we see all that? Well when greet was called, the original greet function wasn't get called (not directly). Remember, our greet function was replaced by the return value from calling our decorator. And that return value is another function (the wrapper function that we defined inside our decorator).

How do decorated functions work?

We defined this function (called wrapper) in our decorator:

def wrapper(*args, **kwargs):
    print("Calling with", args, kwargs)
    return_value = func(*args, **kwargs)
    print("Returning", return_value)
    return return_value
return wrapper

When wrapper was called, it accepted any arguments we gave to it.

It then printed those arguments we passed to it (as a tuple and a dictionary because that's how * and ** capture them).

It called our original function (our greet function) with whatever arguments were given to the wrapper function.

Then it printed out the return value it got from calling the original greet function. Our greet function doesn't actually return anything, so it defaulted to the default None return value, which wrapper then printed out and returned.

The wrapper function sandwiches the decorated function

It's common for the inner function returned by a decorator to be named wrapper because it wraps around the function it decorates.

While our wrapper function does replace the original function, it also augments the original function. This replacement function calls the original function while also doing something extra. It accepts any arguments we might pass to it (because it doesn't know what arguments the decorated function should accept) and then it passes those arguments to the decorated function.

That inner function typically:

  1. Does something before calling the original function
  2. Calls the original function
  3. Does something after calling the original function

You can think of the wrapper function as kind of like a sandwich.

There's the top of the sandwich:

    print("Calling with", args, kwargs)

The middle of the sandwich (where we actually call our function):

    return_value = func(*args, **kwargs)

And the bottom of the sandwich:

    print("Returning", return_value)
    return return_value

It's possible that this decorator makes an open-faced sandwich, by doing something first but not after, or after but not first, but it will also do something in addition to calling the decorated function.

When you have return value from decorated function

We've used our decorator log_me on one function (greet), but we can decorate as many functions as we like.

Here we have a function (get_hypotenuse) that accepts two arguments and returns a value:

>>> from math import sqrt
>>>
>>> @log_me
... def get_hypotenuse(a, b):
...     return sqrt(a**2 + b**2)

We applied the log_me decorator to this get_hypotenuse function as we defined it.

What do you think will happen when we call this function with two arguments?

>>> h = get_hypotenuse(3, 4)

What will we see?

Note that we're capturing the return value (h) here.

When we call get_hypotenuse with 3 and 4 we see this:

>>> h = get_hypotenuse(3, 4)
Calling with (3, 4) {}
Returning 5.0

We see the decorator printed out the positional arguments ((3, 4)) and the keyword arguments ({}):

And then it printed out the return value (5.0).

So if we look at h we should see the same return value:

>>> h
5.0

And we do!

Summary

So a decorator is a function that accepts a function and returns a function (a function decorator is; class decorators work a little differently).

Typically, the function that a decorator returns wraps around our original function. That wrapper function often looks like a sandwich:

  1. Top: it does something first
  2. Middle: then it calls our original function
  3. Bottom: then it does something afterward

That's how to make a decorator in Python.