An Introduction to Decorators: Enhancing Functions in Python

An Introduction to Decorators: Enhancing Functions in Python

Functions are essential parts of programming in Python, and the language provides many patterns that can enhance your experience when using functions.

If you are a functional programming-oriented engineer, the concepts of decorators might interest you. In this article, I will introduce decorators and demonstrate a real-world example of using them.

If you are interested in more content covering topics like the one you're reading, subscribe to my newsletter for regular updates on software programming, architecture, and tech-related insights.

Revisiting Functions

In Python, everything you interact with is an object, including functions.

Here’s a basic function definition in Python:

def func():
    do_some_thing

# or

def func(*args):
    do_something_with_args

A Python function can take arguments or not and return a response. An interesting fact about functions in Python is that you can pass any objects you want, even functions. Check this example:

# Define a function that takes another function as an argument
def apply_function(func, value):
    return func(value)

# Define a couple of simple functions to pass as arguments
def square(x):
    return x * x

def cube(x):
    return x * x * x

# Use the apply_function to apply different functions to a value
result_square = apply_function(square, 4)
result_cube = apply_function(cube, 3)

print(f"Square of 4: {result_square}")  # Output: Square of 4: 16
print(f"Cube of 3: {result_cube}")      # Output: Cube of 3: 27

The apply_function function takes an argument func and a value. It applies the function func to the value value and returns the result. You can see this with result_square and result_cube.

Now, let's go a bit further. Suppose you want to log the execution of the square and cube functions without writing print statements inside these functions. We can achieve this with a decorator.

Introducing Decorators

We will rename apply_function to logger and add an inner function called wrapper to the logger function.

# Define the logger function
def logger(func):
    def wrapper(value):
        print(f"Calling function {func.__name__} with argument {value}")
        result = func(value)
        print(f"Function {func.__name__} returned {result}")
        return result
    return wrapper

Here, we create a decorator. This function takes another function as an argument, and adds logging before and after calling the func, and then returns the wrapper.

In Python, a decorator is a function that takes another function as an argument and extends or alters its behavior. In the example above, we alter the behavior by printing messages before and after the function’s execution.

# Define a couple of simple functions to be decorated
def square(x):
    return x * x

def cube(x):
    return x * x * x

# Manually apply the logger decorator to the functions
logged_square = logger(square)
logged_cube = logger(cube)

# Call the decorated functions
result_square = logged_square(4)
result_cube = logged_cube(3)

print(f"Logged square of 4: {result_square}")  # Output: Logged square of 4: 16
print(f"Logged cube of 3: {result_cube}")      # Output: Logged cube of 3: 27

Here, we manually apply the logger decorator. The logger function is called with square and cube as arguments, creating new decorated versions of these functions (logged_square and logged_cube). When the decorated functions are called, the wrapper function logs the calls and results.

However, the syntax of applying decorators is a little bit long and verbose. Python provides a much simpler way to achieve this.

Applying the Decorator Using @ Syntax

# Define the logger decorator function
def logger(func):
    def wrapper(value):
        print(f"Calling function {func.__name__} with argument {value}")
        result = func(value)
        print(f"Function {func.__name__} returned {result}")
        return result
    return wrapper

# Define a couple of simple functions to be decorated
@logger
def square(x):
    return x * x

@logger
def cube(x):
    return x * x * x

# Call the decorated functions
result_square = square(4)
result_cube = cube(3)

print(f"Logged square of 4: {result_square}")  # Output: Logged square of 4: 16
print f"Logged cube of 3: {result_cube}")      # Output: Logged cube of 3: 27

By using the @ syntax, the code becomes more readable and makes it clear that the square and cube functions are being decorated with the logger decorator.

Now that we understand decorators in Python, let’s explore some real-world examples.

Real-World Examples of Using Decorators

In the previous section, we explored decorators in Python, and how to declare and apply them. Now, let's explore two real-world examples of decorators in action.

Example 1: Counting Database Requests

In this example, we simulate requests made to a database and track the number of requests.

from functools import wraps

request_count = 0

def count_requests(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        global request_count
        request_count += 1
        print(f"Number of requests: {request_count}")
        return func(*args, **kwargs)
    return wrapper

@count_requests
def query_database(query):
    # Simulate a database query
    return f"Results for query: {query}"

print(query_database("SELECT * FROM users"))
print(query_database("SELECT * FROM orders"))
print(query_database("SELECT * FROM products"))

Here, we declare a function called count_requests that takes a function as an argument and contains a wrapper function. The @wraps decorator from functools preserves the metadata of the original function when it is wrapped.

The wrapper function increments the request_count and calls the original function. The query_database function is decorated with @count_requests, which counts and prints the number of times it is called.

Without wraps

If you decide not to use wraps, you can manually copy the attributes of the original function to the wrapper function. This approach is more cumbersome and error-prone.

request_count = 0

def count_requests(func):
    def wrapper(*args, **kwargs):
        global request_count
        request_count += 1
        print(f"Number of requests: {request_count}")
        return func(*args, **kwargs)
    # Manually preserve the original function's metadata
    wrapper.__name__ = func.__name__
    wrapper.__doc__ = func.__doc__
    return wrapper

@count_requests
def query_database(query):
    """Simulate a database query"""
    return f"Results for query: {query}"

print(query_database("SELECT * FROM users"))
print(query_database("SELECT * FROM orders"))
print(query_database("SELECT * FROM products"))

# Verify that metadata is preserved
print(query_database.__name__)  # Output: query_database
print(query_database.__doc__)   # Output: Simulate a database query

This example showcases an interesting fact about decorators: they can take arguments, and metadata can be easily lost if you do not use the wraps built-in decorator.

Example 2: Checking User Authentication and Permissions

In this example, we ensure that the user is connected and has the correct permissions before executing a function.

from functools import wraps

# Example user data
current_user = {
    'is_connected': True,
    'permission_group': 'admin'
}

def check_user_permission(permission_group):
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            if not current_user['is_connected']:
                raise PermissionError("User is not connected")
            if current_user['permission_group'] != permission_group:
                raise PermissionError(f"User does not have {permission_group} permission")
            return func(*args, **kwargs)
        return wrapper
    return decorator

@check_user_permission('admin')
def perform_admin_task():
    return "Admin task performed"

try:
    print(perform_admin_task())
except PermissionError as e:
    print(e)

Here, we declare a function check_user_permission that takes a permission_group argument and returns a decorator. The wrapper function checks if the user is connected and has the required permissions before executing the original function. If the conditions are not met, it raises a PermissionError.

The perform_admin_task function is decorated with @check_user_permission('admin'), ensuring that only connected users with admin permissions can execute it.

In short, the outer function check_user_permission takes the required permission group as an argument and returns the actual decorator function. The inner function decorator takes the original function (func) and defines the wrapper function, which adds permission checks. The wrapper function performs the checks, calls the original function if the user passes, and raises a PermissionError if any check fails.

Now, you know how decorators work in Python. For more information, I recommend reading through the official documentation as it contains everything you might need.

Conclusion

Decorators in Python provide a powerful way to extend or alter the behavior of functions. They are particularly useful for logging, counting function calls, checking permissions, authentification, and much more. By understanding and using decorators, you can write more modular, reusable, and maintainable code.

If you enjoyed this article and want to stay updated with more content, subscribe to my newsletter. I send out a weekly or bi-weekly digest of articles, tips, and exclusive content that you won't want to miss 🚀