Decorators

Objectives

By the end of this chapter, you should be able to:

• Explain what a decorator is
• Create your own decorators to add functionality to existing functions
• Write your own context manager, either as a class or with @contextmanager

💡 Why this matters: Decorators are how libraries and frameworks let you add behavior (logging, timing, access control) to a function without touching its body. You’ve likely already used @staticmethod and @classmethod; now you’ll see how they actually work.

Functions Are Just Objects

A decorator is a function that “decorates,” or enhances, another function. This works because everything in Python is an object, including functions, so you can pass one function into another just like you’d pass a number or a string:

def shout():
    return "WHOA!"

def whisper():
    return "Shhhh"

def perform_action(func):
    print("something is happening")
    return func()

print(perform_action(shout))
# something is happening
# WHOA!

print(perform_action(whisper))
# something is happening
# Shhhh

perform_action takes a function as its argument, prints a message, then calls that function and returns whatever it returns. Nothing here is generator or decorator magic yet, it’s just a function accepting another function as a normal argument.

Building a Decorator by Hand

A decorator is a function that takes another function, wraps it in a new function that adds behavior, and returns that wrapper:

def new_decorator(func):
    def wrap_func():
        print("code before func!")
        func()
        print("code after func!")
    return wrap_func

def decorate_me():
    print("decorate me!")

decorate_me = new_decorator(decorate_me)
decorate_me()
# code before func!
# decorate me!
# code after func!

Walk through what actually happens: new_decorator(decorate_me) runs immediately and returns wrap_func, a brand new function that has decorate_me (the original function) tucked inside it. Reassigning decorate_me = new_decorator(decorate_me) replaces the name decorate_me with that wrapper. So when you call decorate_me() afterward, you’re actually calling wrap_func(), which prints a line, calls the original function you passed in, then prints another line.

The @ Syntax

Reassigning a function to its own decorated version, as above, works, but Python gives you cleaner syntax for it. Prefix the decorating function with @ directly above the function you’re decorating:

def new_decorator(func):
    def wrap_func():
        print("code before func!")
        func()
        print("code after func!")
    return wrap_func

@new_decorator
def decorate_me():
    print("decorate me!")

decorate_me()
# code before func!
# decorate me!
# code after func!

This produces exactly the same output as the manual reassignment above. @new_decorator just applies new_decorator to decorate_me immediately, at the moment decorate_me is defined.

Revisiting the earlier shout/whisper example with decorator syntax:

def perform_action(func):
    def wrap_func():
        print("something is happening")
        return func()
    return wrap_func

@perform_action
def whisper():
    return "Shhhh"

@perform_action
def shout():
    return "WHOA!"

print(whisper())
# something is happening
# Shhhh

print(shout())
# something is happening
# WHOA!

Preserving Function Identity with functools.wraps

There’s a catch: once wrapped, a function’s __name__ and __doc__ point to the wrapper, not the original:

def perform_action(func):
    def wrap_func():
        print("something is happening")
        return func()
    return wrap_func

@perform_action
def shout():
    return "WHOA!"

print(shout.__name__)  # 'wrap_func', not what you'd expect!

The wraps decorator from the functools module fixes this by copying over the original function’s metadata:

from functools import wraps

def perform_action(func):
    @wraps(func)
    def wrap_func():
        print("something is happening")
        return func()
    return wrap_func

@perform_action
def shout():
    return "WHOA!"

print(shout.__name__)  # 'shout', much better

Get in the habit of adding @wraps(func) to every decorator you write. It costs nothing and avoids a confusing bug later, especially if other code inspects a function’s __name__.

Writing Your Own Context Manager

Back in the file I/O module, with open(...) as file: automatically closed the file for you, even if an error happened inside the block. That’s a context manager, and now that you understand decorators, you can build your own.

There are two ways to write one. The first is a class implementing __enter__ and __exit__:

import time

class Timer:
    def __enter__(self):
        self.start = time.time()
        print("timer started")
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        print(f"took {time.time() - self.start:.4f} seconds")

with Timer():
    total = sum(range(1_000_000))
# timer started
# took 0.0123 seconds  (the exact number will vary)

__enter__ runs when the with block starts, and __exit__ runs when it ends, whether the block finished normally or raised an error. That’s the same guarantee finally gives you, wired into the with syntax. Here’s proof that __exit__ really does run even when the block raises an exception:

class Announce:
    def __enter__(self):
        print("entering")
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        print("exiting")

with Announce():
    print("inside the block")
    raise ValueError("something went wrong")
# entering
# inside the block
# exiting
# Traceback (most recent call last):
#     ...
# ValueError: something went wrong

"exiting" still prints before the error propagates and crashes the program. __exit__ cleaned up first, exactly like a finally block would have.

The second way is shorter: the @contextmanager decorator from contextlib turns a generator function into a context manager, using yield to mark the split between “setup” and “teardown”:

from contextlib import contextmanager
import time

@contextmanager
def timer():
    start = time.time()
    print("timer started")
    yield
    print(f"took {time.time() - start:.4f} seconds")

with timer():
    total = sum(range(1_000_000))
# timer started
# took 0.0123 seconds  (the exact number will vary)

Everything before yield runs on entry, and everything after it runs on exit, including if the code inside the with block raises an error. For a simple case like this one, @contextmanager is usually the less verbose option of the two.

Try It

1. Write a decorator that prints how long a function took to run (hint: the time module’s time.time() gives you a timestamp before and after calling the function).
2. Apply your decorator to two different functions using @ syntax.
3. Check __name__ on a decorated function before and after adding @wraps(func).
4. Write the same context manager two ways: once as a class with __enter__/__exit__, once with @contextmanager, and confirm both behave the same in a with block.

Keep this lesson in mind: in the exercises coming up, you’ll revisit the once function from the functions module and rewrite it as a proper decorator.

Recap

• A decorator is a function that wraps another function to add behavior, without modifying the original function’s code.
• @decorator_name above a function definition is shorthand for reassigning the function to its decorated version.
• Without functools.wraps, a decorated function’s __name__ and __doc__ get overwritten by the wrapper’s; @wraps(func) fixes that.
• A context manager runs setup on entry and teardown on exit, even if an error occurs. Write one as a class with __enter__/__exit__, or more concisely with @contextmanager and a single yield.

Next lesson: lambdas and working with dates.