装饰器函数

装饰器函数

装饰器函数旨在不修改原代码的基础上赋予额外的功能

无返回值

def is_prime(num):
   if num < 2:
       return False
   elif num == 2:
       return True
   else:
       for i in range(2, num):
           if num % i == 0:
               return False
       return True

这是一个判断是否为素数的函数，现在想要统计一下执行的时间。

from time import time

def display_time(fun):
	def wrapper():
		t1 = time()
		fun()
		t2 = time
		print(t2 - t1)
	return wrapper

然后在原函数用@调用这个计时函数

@display_time
def count_prime_nums():
    for i in range(2, 10000):
        if is_prime(i):
            print(i)

有返回值

如果原函数有返回值，比如我想统计2~10000之间一共有多少个素数，可以这样写：

def display_time_ret(func):
    def wrapper():
        t1 = time()
        result = func()
        t2 = time()
        print(t2 - t1)
        return result
    return wrapper

@display_time_ret
def count_prime_nums():
    count = 0
    for i in range(2, 10000):
        if is_prime(i):
            count += 1
    return count

有参数

更进一步，如果我还想原函数可以传参，可以这样写：

def display_time_arg(func):
    def wrapper(*args):
        t1 = time()
        result = func(*args)
        t2 = time()
        print(t2 - t1)
        return result
    return wrapper

@display_time_arg
def prime_nums_arg(max_num):
    count = 0
    for i in range(2, max_num):
        if is_prime(i):
            count += 1
    return count

类装饰器

除上述以外还有类装饰器的写法，类装饰器本质上和函数装饰器原理、作用相同，都是为其它函数增加额外的功能。但是相比于函数装饰器，类装饰器具有灵活度大、高内聚、封装性等优点。
一个对象是否可调用，要看它是否实现了__call__方法。类的对象是不能被调用的，因为其中没有实现__call__，如果我们在其中定义一个__call__方法，那么这个类的对象将会变的可调用。

class A:	# 没有实现__call__方法
    pass

a = A()
a()

执行报错：
TypeError: ‘A’ object is not callable

class A:

    def __init__(self):
        self.num = 10

    def __call__(self):
        print("现在实现了__call__方法:", self.num)

a = A()
a()

使用类装饰器可以直接依靠类内部的__call__方法来实现，当使用 @ 形式将类装饰器附加到函数上时，就会调用类装饰器的__call__方法。而不需要向函数装饰器那样，在装饰器函数中定义嵌套函数，来实现装饰功能。

class display_time_cls:

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

    def __call__(self):
        t1 = time()
        self.fun()
        t2 = time()
        print("花费了%.2fs" % (t2-t1))


@display_time_cls
def count_prime_nums():
    count = 0
    for i in range(2, 10000):
        if is_prime(i):
            count += 1
    print("共有%s个素数" % count)

上面的@display_time_cls相当于count_prime_nums= display_time_cls(count_prime_nums)，等号左边的count_prime_nums是类display_time_cls的实例对象，等号右边的count_prime_nums是类装饰器所装饰的函数名count_prime_nums作为参数传递给类display_time_cls。下面执行count_prime_nums()即为调用类display_time_cls的对象count_prime_nums，此时会自动调用类中定义的__call__方法。