Python小技巧:装饰器(Decorator)

装饰器：decorator

今天看了一个关于python中decorator的讲解，第一次接触这个概念，写一篇文章来捋捋自己的理解。

talk is cheap，show me the code：

# 这段代码是一个最简单版本的判断质数的方法，并且输出所有质数和所花费的时间
# 判断某个返回内质数可以有更加简单的方式，使用一个一维数组来表示各个数是否为质数，
# 从2开始，2*2,2*3，2*4...然后3*3,3*4,3*5...，然后n*(n+1),n*(n+2)...
# n的范围为nums的平方根，然后遍历一遍一维数组
# 输出所有数组状态为true的值，即为所有的质数。
import time

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

def prime_nums():
    t1 = time.time()
   for i in range(2,10000):
      if is_prime(i):
         print(i)
    t2 = time.time()
    print(t2-t1)

prime_nums()

使用decorator将计时部分和逻辑部分分离开：

在上面这段代码中，不难发现，如果存在多个方法中需要用到时间记录，就需要写很多重复的t1，t2，t2-t1的操作，这显得很不简洁，所以我们想要将计时部分和逻辑部分分离开来。这里就需要使用装饰器（decorator）：

import time

# 构建decorator，参数就是我们想要运行的函数
def display_time(func):
   # 再定一个wrapper方法，wrapper表示运行func时需要运行哪些内容；
   # 在这个例子中我们需要先取开始时间，然后调用func函数，然后再取结束时间，并输出
   def wrapper():
      t1 = time.time()
      func()
      t2 = time.time()
      print(t2-t1)
   return wrapper

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

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

prime_nums()

在创建好装饰器后，想要调用装饰器，则需要在prime_nums()方法上加一个@display_time，这样当调用prime_nums()时，会去调用display_time(func)方法，并将prime_nums()函数传入display_time(func)的参数func中，然后会调用warpper，按照里面定义的执行顺序进行运行。

存在返回值的decorator设置：

这里的prime_nums()时没有参数返回的，如果有返回该，即现在prime_nums()的目标是统计存在多少个质数，并返回。如何更改decorator呢？

import time

# func中存在返回值，func中return值不会在wrapper中返回
# 需要在wrapper中获取func的返回值，然后手动return这个值
def display_time_return(func):
   # 运行func时需要运行哪些内容
   def wrapper():
      t1 = time.time()
      result = func()
      t2 = time.time()
      print(t2-t1)
      return result
   return wrapper

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

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

count_1 = count_prime_nums()
print(count_1)

以上code就是在func存在返回值时，decorator的定义情况。在调用func时，由于其存在返回值，但是wrapper中并没有将这个返回值进行保留，所以我们需要在wrapper中手动获取这个返回值，即用一个变量result = func()，然后在warpper中返回result，这样在执行count_prime_nums()方法，调用完decorator，会返回func()原本会返回的值。这样就解决了方法中存在返回值的情况。

存在参数传递的decorator设置：

如果方法存在参数怎么办？当前是固定了10000这个值，但是我们想用参数的形式来传入这个值，这样我们想求多少以内的质数个数只需要更改传入参数就可以了。

import time

def display_time_return_args(func):
   # 运行func时需要运行哪些内容
   def wrapper(*args):
      t1 = time.time()
      result = func(*args)
      t2 = time.time()
      print(t2-t1)
      return result
   return wrapper

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

@display_time_return_args
def count_prime_nums_args(maxnum):
   count = 0
   for i in range(2,maxnum):
      if is_prime(i):
         count = count+1
   return count

count_2 = count_prime_nums_args(20000)
print(count_2)

上面的代码就是count_prime_nums_args(arg)存在参数时使用decorator的方式，在wrapper(*args)加上*args表示传入的参数，这个参数可以是一个或者多个，当执行count_prime_nums_args时调用装饰器，方法中的参数会传入decorator中的*args，在func(*args)也加上这些参数的表示，这样就能将有参方法的参数传入到warpper中。

装饰器可以有效的解决逻辑部分和计时部分(或其他辅助功能)进行分离，这样做就可以使得代码更加简洁，无需重复使用相同的一段代码。