小谈python装饰器及numba的基本使用

1. 预热知识

要理解python中的装饰器，就要明白在python中，函数是一种特殊类型的变量，可以作为参数传递给函数，也可以作为返回值返回。比如下面的代码，就是 str_1 作为参数传递给 str_2 ,然后再 str_2 中调用传入的函数。

def str_1():
    print('good day')
    
def str_2(func):
    func()
    
str_2(str_1)

再看下面的这段代码：

def str_2():
    
  def str_1():
      print("abc")
      
  return str_1

str_2()()

这段代码在 str_2 内部定义的函数 str_1 作为一个变量返回，然后我们在调用 str_2时获取到返回之后，继续调用 str_1。

2. 了解 python 中的装饰器

2.1 装饰器的作用

装饰器本质是一个 python 函数，让其他函数在不需要任何代码变动的前提下增加额外的功能，其返回值也是一个函数对象。

2.2 功能实现过程

实例：

现在我们有以下一个函数：

def sum_1(a,b):
    sum_1 = a + b
    return sum_1

我们希望增加一个获得计算函数运行时间的功能，于是将代码修改如下：

import time

def sum_1(a,b):
    time_1 = time.time()
    sum_0 = a + b
    print(time.time()-time_1)
    return sum_0

如果我们有若干个函数需要相同的功能，所以你复制黏贴了代码如下：

import time

def sum_1(a,b):
    time_1 = time.time()
    sum_0 = a + b
    print(time.time()-time_1)
    return sum_0

def sum_2(a,b):
    time_1 = time.time()
    sum_0 = a + b
    print(time.time()-time_1)
    return sum_0

然而你抿了一口手中的咖啡，觉得这样有点low，于是你一拍脑袋，修改如下：

import time

def count_time(func,a,b):
    time_1 = time.time()
    sum_0 = func(a,b)
    print(time.time()-time_1)
    return sum_0

def sum_1(a,b):
    sum_0 = a + b
    return sum_0

def sum_2(a,b):
    sum_0 = a + b
    return sum_0

count_time(sum_1,1,2)
count_time(sum_2,3,4)

这看上去似乎有点模样了，然后我们还是没法直接运行 sum_1(a,b) 获得我们想要的结果，那么该如何不改变函数的调用方式，获得返回的结果呢？你再次抿了口咖啡，飞快地敲下如下的代码：

import time

def count_time(func):
    def count(a,b):
        time_1 = time.time()
        sum_0 = func(a,b)
        print(time.time()-time_1)
        return sum_0
    return count

def sum_1(a,b):
    sum_0 = a + b
    return sum_0

def sum_2(a,b):
    sum_0 = a + b
    return sum_0

sum_1 = count_time(sum_1)
sum_2 = count_time(sum_2)

sum_1(1,2)
sum_2(3,4)

emmmm,这样就舒服多了嘛，你的嘴角露出了一丝弧度，你不禁为自己的机智点了个赞，然后你不禁再次问自己，还能再简化下吗？

2.3 闪亮登场的 Decorator

你冥思苦想一番，用到了语法糖——装饰器。将上面的代码修改如下：

import time

def count_time(func):
    def count(a,b):
        time_1 = time.time()
        sum_0 = func(a,b)
        print(time.time()-time_1)
        return sum_0
    return count

@count_time
def sum_1(a,b):
    sum_0 = a + b
    return sum_0

@count_time
def sum_2(a,b):
    sum_0 = a + b
    return sum_0

sum_1(1,2)
sum_2(3,4)

@符号是一个语法糖，将被包裹的函数作为一个变量传递给装饰函数/类，讲装饰函数/类返回的值替换原本的函数，即替换了以下过程：

sum_1 = count_time(sum_1)
sum_2 = count_time(sum_2)

到这里，你大概对 python 的装饰器有了个初步的了解了吧！

3. numba的基本使用

python是一种动态语言，如果能够让它静态一点，速度会好很多，于是有了 cpython。然后 cpython 还是有诸多不便。于是 numba 就成了一个强大而又方便的替代品。它对 for 循环有很好的效果，实例如下：

from numba import jit
from numpy import arange
import time

@jit
def sum_1(arr):
    M, N = arr.shape
    result = 0.0
    for i in range(M):
        for j in range(N):
            result += arr[i,j]
    return result

a = arange(9999999).reshape(3333333,3)        

start = time.time()
print(sum_1(a))
stop = time.time()
print(stop-start)

使用@jit 后的运行时间大约是0.06s，不使用@jit 的运行时间大约是3.2s，效果提高了50多倍。