Python Decorators: Learn to Create and Use Them

Introduction

In this lab, you will learn about Python Decorators, a powerful feature that provides a concise and reusable way to extend the functionality of a function or a class without permanently modifying it.

You will start by building a simple decorator, then learn how to handle function parameters, apply best practices using functools.wraps, create decorators that accept their own arguments, and finally, implement decorators using classes to manage state.

Bare bone decorator

In this step, you will learn the fundamental structure of a Python decorator. A decorator, in its simplest form, is a function that takes another function as an argument, adds some functionality, and then returns another function.

Note: If you are not familiar with using WebIDE (VS Code like), please learn it first at this lab.

First, open the file step1.py from the file explorer on the left.

Now, copy the following code into the step1.py file. This code defines a simple decorator and applies it to a function.

## Decorator: a function that takes another function and returns a wrapper
def your_decorator(func):
  def wrapper():
    ## Do stuff before func...
    print("Before func!")
    func()  ## Call the original function
    ## Do stuff after func...
    print("After func!")
  return wrapper  ## Return the wrapper function

## @your_decorator is syntactic sugar for: foo = your_decorator(foo)
@your_decorator
def foo():
  print("Hello World!")

foo()  ## Calls wrapper, which calls foo with extra behavior

Let's break down the code:

your_decorator(func): This is the decorator function. It takes one argument, func, which is the function to be decorated.
wrapper(): This is an inner function that "wraps" the original function. It contains the new behavior you want to add. Here, it prints messages before and after calling the original function.
@your_decorator: This is the syntactic sugar for applying the decorator. It's equivalent to writing foo = your_decorator(foo) after the function definition.

Now, run the script in the terminal to see the decorator in action.

python3 step1.py

You should see the following output, which shows that the wrapper function was executed, adding behavior around the original foo function.

Before func!
Hello World!
After func!

Decorator for a function with parameters

In this step, you will extend your decorator to work with functions that accept arguments. A real-world function often has parameters, and our decorator needs to be able to handle them correctly.

To do this, you can use *args and **kwargs in the wrapper function's signature. This allows the wrapper to accept any number of positional and keyword arguments and pass them on to the original function.

Open the file step2.py and add the following code:

## Decorator that works with functions that have parameters
def your_decorator(func):
  def wrapper(*args,**kwargs):  ## Accept any arguments
    ## Do stuff before func...
    print("Before func!")
    func(*args,**kwargs)  ## Pass arguments to original function
    ## Do stuff after func...
    print("After func!")
  return wrapper

@your_decorator
def foo(bar):
  print("My name is " + bar)

foo("Jack")  ## Arguments are passed through wrapper

In this code, the wrapper function is defined with *args and **kwargs. When foo("Jack") is called, the argument "Jack" is captured by *args in the wrapper and then passed to the original func when func(*args, **kwargs) is executed.

Run the script to see the result.

python3 step2.py

The output shows that the argument was successfully passed through the decorator to the original function.

Before func!
My name is Jack
After func!

Template for a basic decorator

In this step, you will discover the best-practice template for creating robust decorators. While the previous decorators work, they have a side effect: they lose the metadata of the original function (like its name __name__ and docstring __doc__).

To solve this, Python's functools module provides a helper decorator called wraps. Additionally, a good decorator should be able to handle functions that return values.

This template is useful for most decorator use-cases. It is valid for functions with or without parameters, and with or without a return value.

Open the file step3.py and add the following code.

import functools

## Best practice decorator template: preserves function metadata and return value
def your_decorator(func):
  @functools.wraps(func)  ## Preserves function name, docstring, etc.
  def wrapper(*args,**kwargs):
    print("Before function execution...")
    result = func(*args,**kwargs)  ## Call function and capture return value
    print("After function execution...")
    return result  ## Return the original function's return value
  return wrapper

@your_decorator
def add(a, b):
  """This function adds two numbers."""
  return a + b

sum_result = add(5, 3)
print(f"Function name: {add.__name__}")
print(f"Docstring: {add.__doc__}")
print(f"Result: {sum_result}")

Key improvements in this template:

import functools: Imports the necessary module.
@functools.wraps(func): This is a decorator applied to the wrapper function. It copies the metadata from the original func (like __name__ and __doc__) to the wrapper function.
result = func(*args, **kwargs): The return value of the original function is captured in the result variable.
return result: The wrapper function returns the captured result, so the decorated function behaves as expected.

Execute the script:

python3 step3.py

The output demonstrates that the function's name and docstring are preserved, and the return value is correctly handled.

Before function execution...
After function execution...
Function name: add
Docstring: This function adds two numbers.
Result: 8

Decorator with parameters

In this step, you will learn how to create decorators that can be customized with their own parameters. This allows for more flexible and configurable decorator usage, such as @your_decorator(arg='value').

To achieve this, you need to create a "decorator factory"—a function that takes arguments and returns a decorator. This adds an extra layer of nesting to our structure.

Let's create a decorator that takes a prefix string and prints it before and after the function call. Open step4.py and add the following code:

import functools

## Decorator factory: returns a decorator based on parameters
def prefix_decorator(prefix):
  def decorator(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
      ## Use the 'prefix' argument from the factory
      print(f"{prefix} Before function call")
      result = func(*args, **kwargs)
      print(f"{prefix} After function call")
      return result
    return wrapper
  return decorator  ## Return the actual decorator function

## Using the decorator with a parameter
@prefix_decorator(prefix="LOG:")
def say_hello(name):
  print(f"Hello, {name}!")

say_hello("Alice")

Here's how it works:

prefix_decorator(prefix="LOG:") is called first. It's the factory function.
It returns the decorator function, which now has access to the prefix variable ("LOG:") via a closure.
This returned decorator is then applied to say_hello, decorating it.

Run the script to see how the parameter is used.

python3 step4.py

The output shows the "LOG:" prefix, which was passed as an argument to the decorator.

LOG: Before function call
Hello, Alice!
LOG: After function call

Class based decorators

In this step, you will explore how to implement decorators using classes. This approach is particularly useful for maintaining state between function calls. We will cover two patterns: decorating a method within a class and defining a decorator as a class itself.

Open the file step5.py to get started.

Pattern 1: Decorating a Class Method

You can define a decorator as a method inside a class and apply it to other methods of that same class. This is useful for creating specific behaviors tied to the class's context.

Add the following code to step5.py. This example defines a decorator to handle exceptions for a class method.

## Pattern 1: Decorating a method within a class
class DecorateMyMethod:

  ## Static method decorator for methods with only 'self' parameter
  def decorator_for_class_method_with_no_args(method):
    def wrapper_for_class_method(self):  ## Only takes self
      try:
        return method(self)  ## Call original method
      except Exception as e:
        print(f"\nWARNING: An exception occurred: {e}\n")
    return wrapper_for_class_method

  def __init__(self,succeed:bool):
    self.succeed = succeed

  @decorator_for_class_method_with_no_args
  def class_action(self):
    if self.succeed:
      print("You succeeded by choice.")
    else:
      raise Exception("Epic fail of your own creation.")

print("--- Testing Pattern 1 ---")
test_succeed = DecorateMyMethod(True)
test_succeed.class_action()

test_fail = DecorateMyMethod(False)
test_fail.class_action()
print("-" * 25)

Pattern 2: Decorator as a Class

A more common and powerful pattern is to define the decorator itself as a class. This is ideal for when the decorator needs to maintain a state. By implementing the __init__() and __call__() methods, you can make an instance of a class behave like a function.

Let's create a decorator that counts how many times a function has been called. Add the following code to the end of step5.py.

## Pattern 2: Class-based decorator for maintaining state
class CountCallNumber:

  def __init__(self, func):
    self.func = func  ## Store the function to decorate
    self.call_number = 0  ## Initialize call counter

  def __call__(self, *args, **kwargs):  ## Makes instance callable
    self.call_number += 1  ## Increment counter
    print("This is execution number " + str(self.call_number))
    return self.func(*args, **kwargs)  ## Call original function

@CountCallNumber
def say_hi(name):
  print("Hi! My name is " + name)

print("\n--- Testing Pattern 2 ---")
say_hi("Jack")
say_hi("James")

__init__(self, func): This is the constructor. It's called when the decorator is applied (@CountCallNumber). It receives the function to be decorated (say_hi) and stores it.
__call__(self, *args, **kwargs): This method makes the class instance callable. When you call the decorated function (say_hi("Jack")), you are actually invoking this __call__ method. It increments the counter and then calls the original stored function.

Now, run the entire script.

python3 step5.py

You will see the output for both patterns. Notice how the call counter increments in the second part.

--- Testing Pattern 1 ---
You succeeded by choice.

WARNING: An exception occurred: Epic fail of your own creation.

-------------------------

--- Testing Pattern 2 ---
This is execution number 1
Hi! My name is Jack
This is execution number 2
Hi! My name is James

Summary

Congratulations on completing the lab! You have gained a solid understanding of Python decorators and their various forms.

In this lab, you learned how to:

Create a basic decorator and understand the @ syntax.
Design decorators that work with functions accepting any number of arguments using *args and **kwargs.
Use the best-practice template with functools.wraps to preserve function metadata and handle return values.
Build decorator factories that allow your decorators to be customized with parameters.
Implement decorators as classes to maintain state between function calls, such as in a call counter.

Decorators are a key feature for writing clean, reusable, and extensible Python code. You can now apply these patterns to enhance your own projects.