使用timeit測試Python函數的性能

timeit是Python標準庫內置的小工具，能夠快速測試小段代碼的性能。

認識timeit

timeit 函數：

timeit.timeit(stmt, setup,timer, number)

參數說明：

• stmt: statement的縮寫，你要測試的代碼或者語句，純文本，默認值是 "pass"
• setup: 在運行stmt前的配置語句，純文本，默認值也是 "pass"
• timer: 計時器，通常忽略這個參數
• number: stmt執行的次數，默認是1000000，一百萬

repeat 函數：

timeit.repeat(stmt, setup, timer, repeat, number)

是timeit的repeat版，能夠指定重複timeit的次數，默認是3次，而後返回一個數組。

舉一個簡單的例子來講明用法：

import timeit
print(timeit.timeit('output = 10*5')) # 0.014560436829924583
print(timeit.repeat('output = 10*5')) # [0.01492984383367002, 0.01342877489514649, 0.013638464966788888]

嗯，看上去沒毛病，實際上誰也不會去測沒有意義的加減乘除，咱們須要測試本身的代碼。

測試多行代碼

測試多行代碼能夠用分號來鏈接語句。

print(timeit.timeit(stmt='a=10;b=10;sum=a+b'))

也能夠用三引號來寫stmt。

import timeit
import_module = "import random"
testcode = ''' 
def test(): 
    return random.randint(10, 100)
'''
print(timeit.repeat(stmt=testcode, setup=import_module))

可是其實都挺扯的，本身的代碼會那麼簡單？咱們是模塊化編程。

測試模塊中的函數

若是你要測試的函數都在一個模塊裏，能夠這樣寫timeit。

import timeit
import random
import arrow

# 本地函數
def stupid1():
    return random.randint(1, 10)

# 依賴其餘函數
def stupid2():
    return stupid1()

# 依賴其餘包或者模塊
def stupid3():
    return arrow.now()

print(timeit.timeit('stupid1()', setup='from __main__ import stupid1'))
print(timeit.timeit('stupid2()', setup='from __main__ import stupid2'))
print(timeit.timeit('stupid3()', setup='from __main__ import stupid3', number=100))

寫成上面這樣的其實仍是單行的模式。

借用default_timer

timeit自帶的default_timer能夠借來用一下。

import timeit
import random
 
def test(): 
    return random.randint(10, 100)
 
starttime = timeit.default_timer()
print("The start time is :",starttime)
test()
print("The time difference is :", timeit.default_timer() - starttime)

命令行的用法

timeit還支持命令行的調用方式，不過我以爲太累了，不必去嘗試。

C:\pythontest>python -m timeit -s 'text="hello world"'
20000000 loops, best of 5: 13.1 nsec per loop

分享一個案例

2月29那天，我想今年是閏年啊，計算閏年有幾種算法啊？孔乙己說有3種：

def is_leap_year_0(year):
    if year % 4 == 0:
        if year % 100 == 0:
            if year % 400 == 0:
                return True
            else:
                return False
        else:
            return True
    else:
        return False


def is_leap_year_1(year):
    return year % 4 == 0 and (year % 100 != 0 or year % 400 == 0)


def is_leap_year_2(year):
    if year % 400 == 0:
        return True
    if year % 100 == 0:
        return False
    if year % 4 == 0:
        return True
    return False

這三種方法那種最好啊？這個不能一律而論吧，由於要看你的參數是什麼，好比1991年不是閏年，方法0和方法1直接就返回了，但方法2還須要走到最後一個if才知道不是閏年。再好比2020年，方法2直接就返回了，可是方法0和1須要走到最裏層的if才獲得結果。因此咱們須要取樣測試才公平，好比從1900年到2900年，每一個函數都跑10000遍。

timeit就不太方便了，它接受的參數哪能那麼複雜，咱們須要包裝一下。

def perf_test(method):
    years = range(1900, 2900)
    if method == 0:
        for y in years:
            is_leap_year_0(y)

    if method == 1:
        for y in years:
            is_leap_year_1(y)

    if method == 2:
        for y in years:
            is_leap_year_2(y)

print(timeit('perf_test(0)', setup='from __main__ import perf_test', number=10000))
print(timeit('perf_test(1)', setup='from __main__ import perf_test', number=10000))
print(timeit('perf_test(2)', setup='from __main__ import perf_test', number=10000))

大家猜猜看哪一個方法結果最好？你必定想不到。

1.6432780250906944
1.7527272370643914
0.0023122059646993876

其餘的思路

timeit其實仍是太弱了，隨便測測還湊合，若是真要檢查性能問題仍是須要用更專業的手段。好比：

• PyCharm Profiler （Pro版功能）
• cProfile
• pycallgraph
• memory_profiler

有機會咱們下次再說。

關於做者：

Toby Qin, Python 技術愛好者，目前從事測試開發相關工做，轉載請註明原文出處。

歡迎關注個人博客 https://betacat.online，你能夠到個人公衆號中去當吃瓜羣衆。