Create your own Python decorator

Overview

In the article A Deeper Look at Python Decorators, I introduced the concept of Python decorators, demonstrated many cool decorators, and explained how to use them.

In this tutorial I will show you how to write your own decorator. As you'll see, writing your own decorator gives you a lot of control and enables a lot of functionality. Without decorators, these features would require a lot of error-prone and repetitive boilerplate that clutters the code, or require completely external mechanisms (such as code generation).

If you don't know anything about decorators, here's a quick refresher. A decorator is a callable object (a function, method, class, or object that has a call() method) that accepts a callable object as input and returns a callable object as output. Typically, the returned callable performs some operations before and/or after calling the input callable. You can use the @ syntax to apply decorators. Lots of examples coming soon...

Hello World Decorator

Let's start with the "Hello world!" decorator. This decorator will completely replace any decorated callable function with a function that just prints "Hello World!"

def hello_world(f):
    def decorated(*args, **kwargs):
        print 'Hello World!'
    return decorated

That's it. Let's see it in action and then explain the different parts and how they work. Suppose we have the following function, which takes two numbers and prints their product:

def multiply(x, y):
    print x * y

If you call, you will get what you expect:

(6, 7)
42

Let's annotate the multiply function with @hello_world and decorate it with the hello_world decorator.

@hello_world
def multiply(x, y):
    print x * y

Now when you call multiply with any arguments (including the wrong data type or the wrong number of arguments), the result is always printing "Hello World!".

multiply(6, 7)
Hello World!

multiply()
Hello World!

multiply('zzz')
Hello World!

OK. How does it work? The original multiply function is completely replaced by the nested decoration function within the hello_world decorator. If we analyze the structure of the hello_world decorator, then you see that it accepts the input callable f (not used in this simple decorator), which defines a nested A function named decorated that accepts any combination of arguments and keyword arguments (defdecorated(*args, **kwargs)) and finally returns the decorated< /strong> function .

Writing function and method decorators

There is no difference between writing function and method decorators. The definition of decorators is the same. The input callable will be a regular function or bound method.

Let's verify it. This is a decorator that just prints the input callable and type before calling it. This is a typical case where the decorator does something and continues calling the original callable.

def print_callable(f):
    def decorated(*args, **kwargs):
        print f, type(f)
        return f(*args, **kwargs)
    return decorated

Note that the last line calls the input callable in a generic way and returns the result. This decorator is non-intrusive because you can decorate any function or method in a working application and the application will continue to work because the decorated function calls the original function and will only have some side effects before.

Let's see it in action. I'm going to decorate our multiplication function and method.

@print_callable
def multiply(x, y):
    print x * y

class A(object):
    @print_callable
    def foo(self):
        print 'foo() here'

When we call functions and methods, the callable is printed, and then they perform the original task:

multiply(6, 7)
 
42

A().foo()
 
foo() here

Decorator with parameters

Decorators can also accept parameters. This ability to configure decorator operations is very powerful, allowing you to use the same decorator in many contexts.

Suppose you're coding too fast and your boss asks you to slow down because you're making other team members look bad. Let's write a decorator that measures how long a function takes to run, and if it takes less than a certain number of seconds t, it will wait until t seconds expire and then return.

The difference now is that the decorator itself accepts a parameter t to determine the minimum running time, and different functions can be decorated with different minimum running times. Also, you'll notice that when introducing decorator parameters, two levels of nesting are required:

import time

def minimum_runtime(t):
    def decorated(f):
        def wrapper(*args, **kwargs):
            start = time.time()
            result = f(*args, **kwargs)
            runtime = time.time() - start
            if runtime < t:
                time.sleep(t - runtime)
            return result
        return wrapper
    return decorated

Let's open it. The decorator itself - the function minimum_runtime takes a parameter t, which represents the minimum running time of the decorated callable function. The input callable f is "pushed down" to the nested decorator function, and the input callable arguments are "pushed down" to another nested function wrapper .

The actual logic happens inside the wrapper function. Record the start time, call the original callable f with its arguments, and store the result. Then it checks the runtime and if it's less than mint then it sleeps for the rest of the time and then returns.

To test it, I will create several functions that call multiplication and decorate them with different delays.

@minimum_runtime(1)
def slow_multiply(x, y):
    multiply(x, y)
    
@minimum_runtime(3)
def slower_multiply(x, y):
    multiply(x, y)

现在，我将直接调用 multiply 以及较慢的函数并测量时间。

import time

funcs = [multiply, slow_multiply, slower_multiply]
for f in funcs:
    start = time.time()
    f(6, 7)
    print f, time.time() - start

这是输出：

42
 1.59740447998e-05
42
 1.00477004051
42
 3.00489807129

正如您所看到的，原始乘法几乎没有花费任何时间，并且较慢的版本确实根据提供的最小运行时间进行了延迟。

另一个有趣的事实是，执行的装饰函数是包装器，如果您遵循装饰的定义，这是有意义的。但这可能是一个问题，特别是当我们处理堆栈装饰器时。原因是许多装饰器还会检查其输入可调用对象并检查其名称、签名和参数。以下部分将探讨此问题并提供最佳实践建议。

对象装饰器

您还可以使用对象作为装饰器或从装饰器返回对象。唯一的要求是它们有一个 __call__() 方法，因此它们是可调用的。下面是一个基于对象的装饰器的示例，它计算其目标函数被调用的次数：

class Counter(object):
    def __init__(self, f):
        self.f = f
        self.called = 0
    def __call__(self, *args, **kwargs):
        self.called += 1
        return self.f(*args, **kwargs)

这是在行动：

@Counter
def bbb():
    print 'bbb'

bbb()
bbb

bbb()
bbb

bbb()
bbb

print bbb.called
3

在基于函数的装饰器和基于对象的装饰器之间进行选择

这主要是个人喜好问题。嵌套函数和函数闭包提供了对象提供的所有状态管理。有些人对类和对象感觉更自在。

在下一节中，我将讨论行为良好的装饰器，而基于对象的装饰器需要一些额外的工作才能表现良好。

行为良好的装饰器

通用装饰器通常可以堆叠。例如：

@decorator_1
@decorator_2
def foo():
    print 'foo() here'

当堆叠装饰器时，外部装饰器（本例中为decorator_1）将接收内部装饰器（decorator_2）返回的可调用对象。如果decorator_1在某种程度上依赖于原始函数的名称、参数或文档字符串，并且decorator_2是简单实现的，那么decorator_2将看不到原始函数中的正确信息，而只能看到decorator_2返回的可调用信息。

例如，下面是一个装饰器，它验证其目标函数的名称是否全部小写：

def check_lowercase(f):
    def decorated(*args, **kwargs):
        assert f.func_name == f.func_name.lower()
        f(*args, **kwargs)
    return decorated

让我们用它来装饰一个函数：

@check_lowercase
def Foo():
    print 'Foo() here'

调用 Foo() 会产生断言：

In [51]: Foo()
---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
ipython-input-51-bbcd91f35259 in module()
----> 1 Foo()

ipython-input-49-a80988798919 in decorated(*args, **kwargs)
      1 def check_lowercase(f):
      2     def decorated(*args, **kwargs):
----> 3         assert f.func_name == f.func_name.lower()
      4     return decorated

但是，如果我们将 check_lowercase 装饰器堆叠在像 hello_world 这样返回名为“decorated”的嵌套函数的装饰器上，结果会非常不同：

@check_lowercase
@hello_world
def Foo():
    print 'Foo() here'

Foo()
Hello World!    

check_lowercase 装饰器没有引发断言，因为它没有看到函数名称“Foo”。这是一个严重的问题。装饰器的正确行为是尽可能多地保留原始函数的属性。

让我们看看它是如何完成的。现在，我将创建一个 shell 装饰器，它仅调用其输入可调用函数，但保留输入函数中的所有信息：函数名称、其所有属性（如果内部装饰器添加了一些自定义属性）及其文档字符串。

def passthrough(f):
    def decorated(*args, **kwargs):
        f(*args, **kwargs)
    decorated.__name__ = f.__name__
    decorated.__name__ = f.__module__
    decorated.__dict__ = f.__dict__
    decorated.__doc__ = f.__doc__    
    return decorated

现在，堆叠在passthrough装饰器之上的装饰器将像直接装饰目标函数一样工作。

@check_lowercase
@passthrough
def Foo():
    print 'Foo() here'

使用@wraps装饰器

此功能非常有用，以至于标准库在 functools 模块中有一个名为“wraps”的特殊装饰器，可以帮助编写与其他装饰器配合良好的适当装饰器。您只需在装饰器中使用 @wraps(f) 装饰返回的函数即可。看看使用 wraps 时 passthrough 看起来有多简洁：

from functools import wraps

def passthrough(f):
    @wraps(f)
    def decorated(*args, **kwargs):
        f(*args, **kwargs)
    return decorated

我强烈建议始终使用它，除非您的装饰器旨在修改其中一些属性。

编写类装饰器

类装饰器是在 Python 3.0 中引入的。他们对整个班级进行操作。定义类时和创建任何实例之前会调用类装饰器。这使得类装饰器几乎可以修改类的每个方面。通常您会添加或修饰多个方法。

让我们直接跳到一个奇特的示例：假设您有一个名为“AwesomeClass”的类，其中包含一堆公共方法（名称不以下划线开头的方法，例如 init），并且您有一个基于单元测试的测试类，名为“AwesomeClassTest”。 AwesomeClass 不仅很棒，而且非常关键，您要确保如果有人向 AwesomeClass 添加新方法，他们也会向 AwesomeClassTest 添加相应的测试方法。这是 AwesomeClass：

class AwesomeClass:
    def awesome_1(self):
        return 'awesome!'

    def awesome_2(self):
        return 'awesome! awesome!'

这是 AwesomeClassTest：

from unittest import TestCase, main

class AwesomeClassTest(TestCase):
    def test_awesome_1(self):
        r = AwesomeClass().awesome_1()
        self.assertEqual('awesome!', r)
        
    def test_awesome_2(self):
        r = AwesomeClass().awesome_2()
        self.assertEqual('awesome! awesome!', r)

if __name__ == '__main__':        
    main()

现在，如果有人添加带有错误的 awesome_3 方法，测试仍然会通过，因为没有调用 awesome_3 的测试。

如何确保每个公共方法始终都有一个测试方法？好吧，当然，你编写一个类装饰器。 @ensure_tests 类装饰器将装饰 AwesomeClassTest 并确保每个公共方法都有相应的测试方法。

def ensure_tests(cls, target_class):
    test_methods = [m for m in cls.__dict__ if m.startswith('test_')]
    public_methods = [k for k, v in target_class.__dict__.items() 
                      if callable(v) and not k.startswith('_')]
    # Strip 'test_' prefix from test method names
    test_methods = [m[5:] for m in test_methods]
    if set(test_methods) != set(public_methods):
        raise RuntimeError('Test / public methods mismatch!')
    return cls

这看起来不错，但有一个问题。类装饰器只接受一个参数：被装饰的类。 Ensure_tests 装饰器需要两个参数：类和目标类。我找不到一种方法来让类装饰器具有类似于函数装饰器的参数。没有恐惧。 Python 有 functools.partial 函数专门用于这些情况。

@partial(ensure_tests, target_class=AwesomeClass)
class AwesomeClassTest(TestCase):
    def test_awesome_1(self):
        r = AwesomeClass().awesome_1()
        self.assertEqual('awesome!', r)

    def test_awesome_2(self):
        r = AwesomeClass().awesome_2()
        self.assertEqual('awesome! awesome!', r)
        
if __name__ == '__main__':
    main()        

运行测试会成功，因为所有公共方法 awesome_1 和 awesome_2 都有相应的测试方法 test_awesome_1 和 test_awesome_2。

----------------------------------------------------------------------
Ran 2 tests in 0.000s

OK

让我们添加一个没有相应测试的新方法awesome_3，然后再次运行测试。

class AwesomeClass:
    def awesome_1(self):
        return 'awesome!'

    def awesome_2(self):
        return 'awesome! awesome!'

    def awesome_3(self):
        return 'awesome! awesome! awesome!'

再次运行测试会产生以下输出：

python3 a.py
Traceback (most recent call last):
  File "a.py", line 25, in module
    class AwesomeClassTest(TestCase):
  File "a.py", line 21, in ensure_tests
    raise RuntimeError('Test / public methods mismatch!')
RuntimeError: Test / public methods mismatch!

类装饰器检测到不匹配并大声清晰地通知您。

结论

编写 Python 装饰器非常有趣，可以让您以可重用的方式封装大量功能。要充分利用装饰器并以有趣的方式组合它们，您需要了解最佳实践和习惯用法。 Python 3 中的类装饰器通过自定义完整类的行为添加了一个全新的维度。

The above is the detailed content of Create your own Python decorator. For more information, please follow other related articles on the PHP Chinese website!