Improve Documentation - To Convert Existing Decorators and Class Based Decorators

[Spill-Tea]

It would be profoundly helpful to provide a walkthrough, demonstrating how to convert common existing decorator patterns, to use the wrapt package.

Take a simple case, of a decorator that measures the execution time of the function it wraps, which may be implemented as either a function or class decorator. For brevity, I reuse the functional decorator to reduce duplicated code. See the following three typical patterns:

import time
import functools

# Wrapper Method 1: as a function
def timed(function: typing.Callable):
    @functools.wraps(function)
    def wrapper(*args, **kwargs):
        t0 = time.perf_counter_ns()
        result = function(*args, **kwargs)
        t1 = time.perf_counter_ns()
        print(f"Time (ms): {(t1 - t0) / 1e6 :,.3f}")
        return result
    return wrapper


# Wrapper Method 2: as a Class, where instantiation takes the function
class Timed:
    def __init__(self, function: typing.Callable):
        self.function = timed(function)
        # potentially other logic goes here ...

    def __call__(self, *args, **kwargs):
        return self.function(*args, **kwargs)


# Wrapper Method 3: as a Class, where __call__ takes the function
class Delayed:
    def __init__(self, delay):
        self.delay = delay

    def __call__(self, function: typing.Callable):
        time.sleep(self.delay)
        return timed(function)

and each of these decorators may be used as follows...

@timed
def costly_function():
    """I <3 documentation"""
    time.sleep(1)


@Timed
def costly_function2():
    """I <3 documentation"""
    time.sleep(1)


@Delayed(1)
def costly_function3():
    """I <3 documentation"""
    time.sleep(1)

From the current documentation, it shows us how to create decorators in the first case "Method 1" - Function as a decorator, as follows:

# Case 1 : Function as a decorator
# Unclear: Should we use this instance somehow?
@wrapt.decorator
def timed(wrapped, instance, args, kwargs):
    t0 = time.perf_counter_ns()
    result = wrapped(*args, **kwargs)
    t1 = time.perf_counter_ns()
    print(f"Time (ms): {(t1 - t0) / 1e6 :,.3f}")
    return result

however, It isn't entirely clear, how we convert classes as Decorators? From the documentation, we have the following, which is another function:

# from current documentation
def with_arguments(myarg1, myarg2):
    @wrapt.decorator
    def wrapper(wrapped, instance, args, kwargs):
        return wrapped(*args, **kwargs)
    return wrapper


# Case 3: Using the formula above, we can convert our Delayed Class
# into a function based wrapper using wrapt as follows:
def delayed(delay: int):
    @wrapt.decorator
    def inner(wrapped, instance, args, kwargs):
        time.sleep(delay)
        t0 = time.perf_counter_ns()
        result = wrapped(*args, **kwargs)
        t1 = time.perf_counter_ns()
        print(f"Time (ms): {(t1 - t0) / 1e6 :,.3f}")
        return result
    return inner

In these simple cases, we could easily convert our class based method decorators, into functions as shown above, but in more complicated examples (where we need to manage more states of the class), this may not be so easy. Is there a way to create class based wrappers using wrapt?

[GrahamDumpleton]

Here are a whole bunch of ways one can use a class object to hold state for the decorator. Should give you a few ideas.

Key thing is ensuring that introspection still works. You cannot use the same structure as your method 2 uses if you want introspection to yield the correct results. As is, the prototype of __call__ in your example would be what results when introspection is done.

import wrapt
import functools
import inspect

class CBD1:

    def __init__(self, *params):
        self._params = params

    @wrapt.decorator
    def __call__(self, wrapped, instance, args, kwargs):
        try:
            print("before", self._params)
            return wrapped(*args, **kwargs)
        finally:
            print("after")

    # Use same class to encapsulate a family of decorators.

    @wrapt.decorator
    def variant(self, wrapped, instance, args, kwargs):
        try:
            print("before", self._params)
            return wrapped(*args, **kwargs)
        finally:
            print("after")

    # Following could have been applied on __call__() if desired.

    def variant_with_params(self, *params):
        @wrapt.decorator
        def wrapper(wrapped, instance, args, kwargs):
            try:
                print("before", self._params, params)
                return wrapped(*args, **kwargs)
            finally:
                print("after")

        return wrapper

    def variant_with_optional_params(self, wrapped=None, *, param3=None, param4=None):
        if wrapped is None:
            return functools.partial(self.variant_with_optional_params,
                    param3=param3, param4=param4)

        @wrapt.decorator
        def wrapper(wrapped, instance, args, kwargs):
            try:
                print("before", self._params, (param3, param4))
                return wrapped(*args, **kwargs)
            finally:
                print("after")

        return wrapper(wrapped)

# Unique state object per decorated function.

@CBD1("param1", "param2")
def func1():
    print("func1")

print("func1", inspect.signature(func1), sep="")

func1()

@CBD1("param1", "param2").variant
def func2():
    print("func2")

print("func2", inspect.signature(func2), sep="")

func2()

@CBD1("param1", "param2").variant_with_params("param3", "param4")
def func3():
    print("func3")

print("func3", inspect.signature(func3), sep="")

func3()

@CBD1("param1", "param2").variant_with_optional_params
def func4():
    print("func4")

print("func4", inspect.signature(func4), sep="")

func4()

@CBD1("param1", "param2").variant_with_optional_params(param3="param3", param4="param4")
def func5():
    print("func5")

print("func5", inspect.signature(func5), sep="")

func5()

# Shared state object for a group of decorated functions.

cbd1 = CBD1("param1", "param2")

@cbd1
def func6():
    print("func6")

print("func6", inspect.signature(func6), sep="")

func6()

@cbd1.variant
def func7():
    print("func7")

print("func7", inspect.signature(func7), sep="")

func7()

@cbd1.variant_with_params("param3", "param4")
def func8():
    print("func8")

print("func8", inspect.signature(func8), sep="")

func8()

# Using functions to hide state object construction.

def cbd2(wrapped):
    return CBD1()(wrapped)

@cbd2
def func9():
    print("func9")

print("func9", inspect.signature(func9), sep="")

func9()

def cbd3(*params):
    def wrapper(wrapped):
        return CBD1(*params)(wrapped)
    return wrapper

@cbd3("param1", "param2")
def func10():
    print("func10")

print("func10", inspect.signature(func10), sep="")

func10()

def cbd4(*params):
    def wrapper(wrapped):
        return CBD1().variant_with_params(*params)(wrapped)
    return wrapper

@cbd4("param3", "param4")
def func11():
    print("func11")

print("func11", inspect.signature(func11), sep="")

func11()

# Using raw FunctionWrapper object proxy.

class CBD2(wrapt.FunctionWrapper):

    def __init__(self, wrapped):
      super().__init__(wrapped, self.wrapper)

    def wrapper(self, wrapped, instance, args, kwargs):
        try:
            print("before")
            return wrapped(*args, **kwargs)
        finally:
            print("after")

def cbd5(wrapped):
    return CBD2(wrapped)

@cbd5
def func12():
    print("func12")

print("func12", inspect.signature(func12), sep="")

func12()

Not that you must use try/finally pattern if want to run after action when wrapped function raises an exception.