Python: Decorator

ishant_mrinal · March 12, 2023, 4:32am

Python decorators are a powerful feature that allow programmers to modify the behavior of functions or classes in a flexible and reusable way. Decorators are a form of metaprogramming that can be used to add functionality to existing code without modifying it directly.

In this blog post, we’ll explore what decorators are, how they work, and how they can be used to simplify and enhance Python code.

What are decorators?

In Python, a decorator is a function that takes another function as input, modifies it in some way, and returns the modified function as output. Decorators are typically used to modify the behavior of a function without changing its source code. For example, a decorator might add logging or timing functionality to a function, or it might wrap the function in a try-except block to handle exceptions.

Here’s an example of a decorator that adds logging functionality to a function:

def log_decorator(func):
  def wrapper(*args, **kwargs):
    print(f"Calling function {func.__name__} with args {args} and kwargs {kwargs}")
    return func(*args, **kwargs)
  return wrapper

@log_decorator
def my_function(x, y):
  return x + y

In this example, the log_decorator function takes another function as input (func ) and returns a new function (wrapper ) that logs the arguments passed to the original function and then calls the original function. The @log_decorator syntax is a shorthand for applying the log_decorator to the my_function function.

How do decorators work?

When a decorator is applied to a function, the decorator function is called with the original function as its input. The decorator function then returns a new function that replaces the original function. When the new function is called, it calls the original function with any arguments and keyword arguments that were passed to it.

Here’s a more detailed breakdown of how the log_decorator example works:

The log_decorator function takes a function func as input.
The wrapper function is defined inside the log_decorator function. The wrapper function takes any number of positional arguments (*args) and any number of keyword arguments (**kwargs).
The wrapper function logs the name of the original function (func.__name__) and the arguments passed to it (args and kwargs).
The wrapper function then calls the original function (func) with the same arguments and keyword arguments (*args and **kwargs).
The wrapper function returns the result of calling the original function (func).

When the log_decorator function is applied to the my_function function using the @log_decorator syntax, the following steps occur:

The log_decorator function is called with the my_function function as its input.
The wrapper function is returned as the new version of the my_function function.
When the my_function function is called, the wrapper function is called instead.
The wrapper function logs the arguments passed to the my_function function and then calls the original my_function function.
The my_function function returns the result of the computation (x + y).

Functools.wraps

We need functools.wrap for decorators because it helps to preserve important metadata of the original function when creating a new function through decoration.

When you decorate a function with another function, the decorated function replaces the original function in the namespace. However, the decorated function may lose important metadata such as the function name, docstring, and parameter annotations. This can be a problem, especially if you want to use the decorated function in introspection or debugging.

To preserve this metadata, you can use the @wraps decorator from the functools module. This decorator is used to wrap the decorated function with a new function that carries over the metadata of the original function. This way, the decorated function retains the same name, docstring, and parameter annotations as the original function, which makes it easier to work with.

In addition to preserving metadata, @wraps also provides other benefits, such as improving the readability of code by providing a clear link between the decorated and original functions, and enabling better error reporting by associating the correct function name and line number with any error messages that may arise.

Overall, the @wraps decorator is an important tool for creating decorators that maintain the same API and behavior as the original function, while also adding new functionality to the function.

Here’s an example that demonstrates the issue with a decorator that doesn’t use functools.wrap to preserve metadata:

def my_decorator(func):
  def wrapper(*args, **kwargs):
    print("Before function call")
    result = func(*args, **kwargs)
    print("After function call")
    return result
  return wrapper

@my_decorator
def my_function(x: int, y: int) -> int:
  """This is my function"""
  return x + y

print(my_function.__name__)  # prints "wrapper"
print(my_function.__doc__)  # prints "None"
print(my_function.__annotations__)  # prints "{}"

Using functools.wraps

import functools


def my_decorator(func):

  @functools.wraps(func)
  def wrapper(*args, **kwargs):
    print("Before function call")
    result = func(*args, **kwargs)
    print("After function call")
    return result

  return wrapper


@my_decorator
def my_function(x: int, y: int) -> int:
  """This is my function"""
  return x + y


print(my_function.__name__) # prints "my_function"
print(my_function.__doc__) # prints "This is my function"
print(my_function.__annotations__) # prints "{'x': <class 'int'>, 'y': <class 'int'>, 'return': <class 'int'>}"

Execution flow of decorators with inputs

When you apply a decorator with inputs to a function, the execution flow is as follows:

The decorator function is called with the inputs provided. The decorator function should return another function that takes the original function as an argument.
The returned function is called with the original function as an argument.
The returned function creates a wrapper function that wraps the original function. This wrapper function can modify the behavior of the original function in some way, for example, by adding additional functionality before or after the original function is called.
The wrapper function is returned by the returned function and is used as the new implementation of the original function.

Here’s an example that demonstrates this execution flow:

import functools


def my_decorator_with_input(arg1, arg2):

  def decorator(func):

    @functools.wraps(func)
    def wrapper(*args, **kwargs):
      print("Before function call")
      result = func(*args, **kwargs)
      print("After function call")
      return result

    return wrapper

  return decorator


@my_decorator_with_input("arg1", "arg2")
def my_function(x, y):
  """
    This is my function
    """
  return x + y


print(my_function(3, 4)) # prints "Before function call", "After function call", and "7"

In this example, the my_decorator_with_input function takes two arguments arg1 and arg2. It returns another function that takes the original function func as an argument. The returned function creates a wrapper function that prints “Before function call” before calling the original function and “After function call” after calling the original function.

When you apply the decorator to the my_function function, the execution flow is as follows:

my_decorator_with_input("arg1", "arg2") is called with the inputs “arg1” and “arg2”. It returns a function that takes my_function as an argument.
my_decorator_with_input("arg1", "arg2")(my_function) is called with my_function as an argument. It returns a wrapper function that wraps my_function.
my_function is replaced by the wrapper function returned by my_decorator_with_input("arg1", "arg2")(my_function).
When my_function(3, 4) is called, the wrapper function is executed. It prints “Before function call”, calls the original my_function with arguments 3 and 4, which returns 7, and then prints “After function call”. The wrapper function then returns 7, which is printed to the console.

How can decorators be used?

Decorators can be used to add functionality to existing code without modifying it directly. This can be particularly useful in situations where modifying the existing code is not possible or desirable. Decorators can also be used to simplify and enhance code by encapsulating complex functionality in a reusable decorator.

Timing function calls with parameters

import time
import functools


def time_this(n):

  def decorator(func):

    @functools.wraps(func)
    def wrapper(*args, **kwargs):
      tic = time.time()
      result = func(*args, **kwargs)
      toc = time.time()
      print(f"{func.__name__} took {round(toc-tic, 3)} seconds with {n}")
      return result

    return wrapper

  return decorator


@time_this(n=10)
def my_function(x, y):
  time.sleep(1)
  return x + y


my_function(1, 2) # prints "my_function took 1.002 seconds with 10" and returns 3

In this example, the time_this decorator takes a parameter n and returns a decorator function that times the decorated function’s execution with that parameter. The wrapper function uses the time module to record the start and end time of the function’s execution, and then prints the function’s name, the execution time, and the value of n . The @wraps decorator is used to preserve the original function’s metadata, such as its name and docstring.

Validating function arguments

import functools


def validate_args_decorator(*types):

  def decorator(func):

    @functools.wraps(func)
    def wrapper(*args):
      for i, arg in enumerate(args):
        if not isinstance(arg, types[i]):
          raise ValueError(f"Argument {i} must be of type {types[i].__name__}")
      return func(*args)

    return wrapper

  return decorator


@validate_args_decorator(int, str)
def my_function(x, y):
  return str(x) + y


print(my_function.__name__)

my_function(1, '2') # returns '12'
my_function('1', '2') # raises ValueError

In this example, the validate_args_decorator function is used to validate the types of the arguments passed to the decorated function. The decorator takes a variable number of type arguments (*types ) and returns a decorator function that takes the original function (func ) as its input. The wrapper function checks the types of each argument and raises a ValueError if any of them are of the wrong type. If all the arguments are of the correct types, it calls the original function with the same arguments.