【Python】函数「四」——给类方法添加装饰器

给类方法添加装饰器

我们有一个Task类，该类有一个类方法run，用于输出任务开始然后调用sleep阻塞一秒：

from time import sleep

class Task:
    def run(self):
        print('任务开始执行')
        sleep(1)
        print('任务执行结束')

task = Task()
task.run()

输出结果：

任务开始执行
任务执行结束

现在，我们需要写一个装饰器为run方法装饰，打印run方法的总耗时，我们可以这么写：

from time import sleep, time

def timer(func):
    def wrapper(*args, **kwargs):
        start = time()
        func(*args, **kwargs)
        print(time()-start)
    return wrapper

class Task:
    @timer
    def run(self):
        print('任务开始执行')
        sleep(1)
        print('任务执行结束')

task = Task()
task.run()

输出结果：

任务开始执行
任务执行结束
1.0101933479309082

如代码所示，和普通函数装饰器一样，我们只需要把参数作为可变参数args和kwargs传入，然后在函数开始之前记录开始时间start，结束后输出结束时间与start的差即可。

那么现在，我们引入一个self.tasks，用于记录已经完成的任务数。此时，如果不希望修改run方法的话，咱们就需要修改装饰器，在run方法结束后将任务完成数self.tasks自增1，改写如下：

from time import sleep, time

def timer(func):
    def wrapper(self, *args, **kwargs):
        start = time()
        func(self, *args, **kwargs)
        print(time()-start)
        self.tasks += 1
    return wrapper

class Task:
    def __init__(self):
        self.tasks = 0

    @timer
    def run(self):
        print('任务开始执行')
        sleep(1)
        print('任务执行结束')
    

task = Task()
print(task.tasks)
task.run()
print(task.tasks)

输出结果：

0
任务开始执行
任务执行结束
1.001990795135498
1

可以看到，完成任务数加1了，说明装饰器没有问题。而在装饰器代码中，因为要使用self.tasks，我们索性把self参数写出来，而且实际上普通类方法默认第一个参数是当前对象self的，所以这样写更符合类方法装饰器的写法，之后大家只需要按照这个demo来为类方法写装饰器即可。

---

那么细心的同学可能会问，既然可以用self调用对象的属性和方法，那可不可以把func(self, *args, **kwargs)改成self.func(*args, **kwargs)呢？

答案是不可以。装饰器是在类定义时就生效的，拿上面的例子来说，不使用语法糖的装饰语句是run = timer(run)，此时传过去的run参数不是当前对象的run，而是类Task类的类方法run，也就是Task.run，所以在self中是不存在这个属性的。因此，在调用的时候采用的是类方法的调用方式而不是对象方法的调用方式，所以是func(self, *args, **kwargs)（在本例中等价于Task.run(self, *args, **kwargs)）。

---

总结

在为普通方法写装饰器时，通用写法是这样：

def decorator(func):
	def wrapper(*args, **kwargs):
		...
		func(*args, **kwargs)
		...
	return wrapper

而普通类方法的特点是第一个参数必为当前对象self，所以通用写法是这样：

def decorator(func):
	def wrapper(self, *args, **kwargs):
		# 可以用self.<attribute>的方式调用当前对象的变量、方法等
        # 比如self.tasks += 1，将任务完成数加1
		...
		func(self, *args, **kwargs)
		...
	return wrapper