There are several different scenarios to which decorators might apply. It can also be the case that we need to use the same decorator for objects that fall into these different multiple scenarios, for instance, if we want to reuse our decorator and apply it to a function, a class, a method, or a static method.

If we create the decorator, just thinking about supporting only the first type of object we want to decorate, we might notice that the same decorator does not work equally well on a different type of object. The typical example is where we create a decorator to be used on a function, and then we want to apply it to a method of a class, only to realize that it does not work. A similar scenario might occur if we designed our decorator for a method, and then we want it to also apply for static methods or class methods.

When designing decorators, we typically think about reusing code, so we will want to use that decorator for functions and methods as well.

Defining our decorators with the signature *args, and **kwargs, will make them work in all cases, because it's the most generic kind of signature that we can have. However, sometimes we might want not to use this, and instead define the decorator wrapping function according to the signature of the original function, mainly because of two reasons:

• It will be more readable since it resembles the original function.
• It actually needs to do something with the arguments, so receiving *args and **kwargs wouldn't be convenient.

Consider the case on which we have many functions in our code base that require a particular object to be created from a parameter. For instance, we pass a string, and initialize a driver object with it, repeatedly. Then we think we can remove the duplication by using a decorator that will take care of converting this parameter accordingly.

In the next example, we pretend that DBDriver is an object that knows how to connect and run operations on a database, but it needs a connection string. The methods we have in our code, are designed to receive a string with the information of the database and require to create an instance of DBDriver always. The idea of the decorator is that it's going to take place of this conversion automatically—the function will continue to receive a string, but the decorator will create a DBDriver and pass it to the function, so internally we can assume that we receive the object we need directly.

An example of using this in a function is shown in the next listing:

import logging
from functools import wraps

logger = logging.getLogger(__name__)


class DBDriver:
    def __init__(self, dbstring):
        self.dbstring = dbstring

    def execute(self, query):
        return f"query {query} at {self.dbstring}"


def inject_db_driver(function):
    """This decorator converts the parameter by creating a ``DBDriver``
    instance from the database dsn string.
    """
    @wraps(function)
    def wrapped(dbstring):
        return function(DBDriver(dbstring))
    return wrapped


@inject_db_driver
def run_query(driver):
    return driver.execute("test_function")

It's easy to verify that if we pass a string to the function, we get the result done by an instance of DBDriver, so the decorator works as expected:

>>> run_query("test_OK")
'query test_function at test_OK'

But now, we want to reuse this same decorator in a class method, where we find the same problem:

class DataHandler:
    @inject_db_driver
    def run_query(self, driver):
        return driver.execute(self.__class__.__name__)

We try to use this decorator, only to realize that it doesn't work:

>>> DataHandler().run_query("test_fails")
Traceback (most recent call last):
  ...
TypeError: wrapped() takes 1 positional argument but 2 were given

What is the problem?

The method in the class is defined with an extra argument—self.

Methods are just a particular kind of function that receives self (the object they're defined upon) as the first parameter.

Therefore, in this case, the decorator (designed to work with only one parameter, named dbstring), will interpret that self is said parameter, and call the method passing the string in the place of self, and nothing in the place for the second parameter, namely the string we are passing.

To fix this issue, we need to create a decorator that will work equally for methods and functions, and we do so by defining this as a decorator object, that also implements the protocol descriptor.

Descriptors are fully explained in Chapter 7, Using Generators, so, for now, we can just take this as a recipe that will make the decorator work.

The solution is to implement the decorator as a class object and make this object a description, by implementing the __get__ method.

from functools import wraps
from types import MethodType


class inject_db_driver:
    """Convert a string to a DBDriver instance and pass this to the 
       wrapped function."""

    def __init__(self, function):
        self.function = function
        wraps(self.function)(self)

    def __call__(self, dbstring):
        return self.function(DBDriver(dbstring))

    def __get__(self, instance, owner):
        if instance is None:
            return self
        return self.__class__(MethodType(self.function, instance))

Details on descriptors will be explained in Chapter 6, Getting More Out of Our Objects with Descriptors, but for the purposes of this example, we can now say that what it does is actually rebinding the callable it's decorating to a method, meaning that it will bind the function to the object, and then recreate the decorator with this new callable.

For functions, it still works, because it won't call the __get__ method at all.