python学习笔记：多种方法实现装饰器及常见误区

本文链接：https://blog.csdn.net/LMFranK/article/details/104969266

文章目录

前言

大部分的文章在介绍装饰器，基本是以函数实现装饰器为主。事实上，装饰器也可以使用其他方式去实现。

正文

本文中将介绍如何实现装饰器，包括使用类、wrapt模块的方式，并且提醒一些常见错误。

1. 闭包

在函数内部再定义一个函数，并且这个函数用到了外边函数的变量，那么将这个函数以及用到的一些变量称之为闭包.

def func(num):

    def inner_func(inner_num):
        print(f"in inner_func, inner_num is {inner_num}")
        return num + inner_num

    return inner_func

res = func(20)
print(res)
print(res(100))

输出：

<function func.<locals>.inner_func at 0x00000139E82B8318>
in inner_func, inner_num is 100
120

使用闭包可以代码复用提高效率，但是由于闭包引用了外部函数的局部变量，因此外部函数的局部变量并不会立刻释放，这会带来额外的开销。

2. 装饰器

2.1 无参装饰器

from time import ctime, sleep

def timer(func):
    def wrapper():
        print(f"{func.__name__} called at {ctime()}")
        func()
    return wrapper

@timer
def fn():
    print("hello, world")

fn()
sleep(2)
fn()

输出：

fn called at Thu Mar 19 15:56:20 2020
hello, world
fn called at Thu Mar 19 15:56:22 2020  # 间隔2s后再次调用函数
hello, world

2.2 被装饰的函数有参数：

from time import ctime, sleep
from functools import reduce

def timer(func):
    def wrapper(*args, **kwargs):
        print(f"{func.__name__} called at {ctime()}")
        print("传入的参数：", *args, **kwargs)
        func(*args, **kwargs)
    return wrapper

@timer
def fn(a, b):
    print(a + b)

fn(3, 5)
sleep(2)
print("==================================")
fn(2, 4)

输出：

fn called at Thu Mar 19 16:39:40 2020
传入的参数： 3 5
8
==================================
fn called at Thu Mar 19 16:39:42 2020
传入的参数： 2 4
6

2.3 装饰器中的return：

from time import ctime, sleep

def timer(func):
    def wrapper():
        print(f"{func.__name__} called at {ctime()}")
        func()
    return wrapper

@timer
def fn():
    print("I am fn")

@timer
def get_info():
    return '----hahah---'

fn()
print("=============================")
sleep(2)
fn()
print("=============================")
print(get_info())

输出：

I am fn
=============================
fn called at Thu Mar 19 16:41:45 2020
I am fn
=============================
get_info called at Thu Mar 19 16:41:45 2020
None

如果将装饰器中修改为return func()，即可以打印出get_info()的return值，因此一般情况下，为了让装饰器更通用，可以加上return。

2.4 带参数的装饰器

from time import ctime, sleep

def timer(pre="hello"):
    def decorator(func):
        def wrapper():
            print(f"{func.__name__} called at {ctime()} {pre}")
            return func()
        return wrapper
    return decorator

@timer()
def fn1():
    print("I am fn1")

@timer("python")
def fn2():
    print("I am fn2")

@timer("world")
def fn3():
    print("I am fn3")

fn1()
print("==============================")
sleep(2)
fn2()
print("==============================")
sleep(2)
fn3()

输出：

fn1 called at Thu Mar 19 16:50:39 2020 hello
I am fn1
==============================
fn2 called at Thu Mar 19 16:50:41 2020 python
I am fn2
==============================
fn3 called at Thu Mar 19 16:50:43 2020 world
I am fn3

3. 类装饰器

绝大多数装饰器的实现都是基于函数，但是这并非是唯一的实现方式。只要是一个callable对象，就可以作为装饰器。

def func():
    pass

print(type(func))  # <class 'function'>
print(callable(func))  # True

显然，函数是callable，那么我们也可以通过__call__这个魔法方法是类变成callable。

class Test(object):

    def __call__(self, *args, **kwargs):
        print("实现了__call__")

test = Test()
print(type(test))  # <class '__main__.Test'>
print(callable(test))  # True
# 调用test
test()  # 实现了__call__

基于这个特性，我们便可以通过类实现装饰器。
下面展示一个类装饰器，用于记录函数的调用次数。

class Counter:

    count = 0

    def __init__(self, func):
        self.func = func

    def __call__(self, *args, **kwargs):
        Counter.count += 1
        return self.func(*args, **kwargs)

@Counter
def fn():
    pass

for i in range(10):
    fn()

print(fn.count)  # 10

相较于纯函数的调用，类装饰器有以下优势：

实现有状态的装饰器时，操作类属性比操作闭包内变量更符合直觉、不易出错
实现为函数扩充接口的装饰器时，使用类包装函数，比直接为函数对象追加属性更易于维护
更容易实现一个同时兼容装饰器与上下文管理器协议的对象（参考 unitest.mock.patch）
当然我们也可以不使用__call__方法去实现类装饰器，下面我们尝试实现一个简单的@property，使得类中的函数以属性的方式调用。

class Decrator(object):
    def __init__(self, func):
        self.func = func

    def __get__(self, instance, owner):
        """
        :param instance: 代表实例，sum中的self
        :param owner: 代表类本身，Test类
        """
        pass
        return self.func(instance)


class Test(object):
    def __init__(self):
        self.result = 0

    @Decrator
    def sum(self):
        print("hello, world")

t = Test()
t.sum  # hello, world

4. wrapt模块

在写装饰器时，往往我们需要面对函数的多层嵌套，可读性不是很高。而且，当我们把原先应用在函数的装饰器，应用到类方法时，会发生错误。
在下面这个例子中，实现了一个生成随机数并注入为函数参数的装饰器。

import random


def provide_number(min_num, max_num):
    """
    随机生成一个在 [min_num, max_num] 范围的整数，追加为函数的第一个位置参数
    """
    def wrapper(func):
        def decorated(*args, **kwargs):
            num = random.randint(min_num, max_num)
            # 将 num 作为第一个参数追加后调用函数
            return func(num, *args, **kwargs)

        return decorated

    return wrapper

@provide_number(1, 100)
def print_random_number(num):
    print(num)

print_random_number()  # 随机生成一个1-100范围的整数

然而，当我们把@provide_number应用到类方法时，得到结果却不是我们想要的。

class Test(object):
    @provide_number(1, 100)
    def print_random_number(self, num):
        print(num)

t = Test()
t.print_random_number()

当调用类方法时，输出的结果是类实例：<__main__.Test object at 0x0000015DCB537E08>。显然，这不符合我们的要求。
wrapt模块便可以解决这个问题。使用wrapt你只需要定义一个装饰器函数，但是函数签名是固定的，必须是(wrapped, instance, args, kwargs)

import random
import wrapt

def provide_number(min_num, max_num):
    """
    随机生成一个在 [min_num, max_num] 范围的整数，追加为函数的第一个位置参数
    """
    @wrapt.decorator
    def wrapper(wrapped, instance, args, kwargs):
        """

        :param wrapped: 被装饰的函数或类方法
        :param instance:
        - 如果被装饰者为普通类方法，该值为类实例
        - 如果被装饰者为 classmethod 类方法，该值为类
        - 如果被装饰者为类/函数/静态方法，该值为None
        :param args: 调用时的位置参数
        :param kwargs: 调用时的关键字参数
        """

        num = random.randint(min_num, max_num)
        args = (num, ) + args
        return wrapped(*args, **kwargs)
    return wrapper

        
class Test(object):
    @provide_number(1, 100)
    def print_random_number(self, num):
        print(num)

t = Test()
t.print_random_number()

常见错误

1. 使用functools.wraps()装饰内层函数

def decorator(func):
    def wrapper(*args, **kwargs):
        """wrapper doc
        """
        print(f"func.__name__是：{func.__name__}")
        return func(*args, **kwargs)
    return wrapper

@decorator
def fn():
    """fn doc
    """
    pass

fn()  # func.__name__是：fn
print(f"fn.__name__是：{fn.__name__}")  # fn.__name__是：wrapper
print(f"fn.__doc__是：{fn.__doc__}")  # fn.__doc__是：wrapper doc

使用装饰器后，被装饰函数的函数签名发生了改变，全部变成了内层函数wrapper的值。这虽然是个小bug，但是有可能会在某些应用场景出问题。此时，我们可以通过functools模块中的wraps解决这个问题。

import functools


def decorator(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        """wrapper doc"""
        print(f"func.__name__是：{func.__name__}")
        return func(*args, **kwargs)
    return wrapper

@decorator
def fn():
    """fn doc"""
    pass

fn()  # func.__name__是：fn
print(f"fn.__name__是：{fn.__name__}")  # fn.__name__是：fn
print(f"fn.__doc__是：{fn.__doc__}")  # fn.__doc__是：fn doc

这样处理后，装饰器就不会影响被装饰函数了。

2. 修改外层变量时记得使用 nonlocal

装饰器是对函数对象的一个高级应用。在编写装饰器的过程中，你会经常碰到内层函数需要修改外层函数变量的情况。就像下面这个装饰器一样：

import functools


def counter(func):
    """装饰器：记录并打印调用次数"""
    count = 0

    @functools.wraps(func)
    def decorated(*args, **kwargs):
        # 次数累加
        count += 1
        print(f"Count: {count}")
        return func(*args, **kwargs)

    return decorated

@counter
def fn():
    pass

fn()

为了统计函数调用次数，我们需要在decorated函数内部修改外层函数定义的count变量的值。但是，上面这段代码是有问题的，在执行它时解释器会报错:

Traceback (most recent call last):
  File "test.py", line 21, in <module>
    fn()
  File "test.py", line 11, in decorated
    count += 1
UnboundLocalError: local variable 'count' referenced before assignment

这个错误是由 counter 与 decorated 函数互相嵌套的作用域引起的。当解释器执行到 count += 1 时，并不知道 count 是一个在外层作用域定义的变量，它把 count 当做一个局部变量，并在当前作用域内查找。最终却没有找到有关 count 变量的任何定义，然后抛出错误。
为了解决这个问题，我们可以使用nonlocal关键字。

@functools.wraps(func)
    def decorated(*args, **kwargs):
        # 次数累加
        nonlocal count
        count += 1
        print(f"Count: {count}")
        return func(*args, **kwargs)