在 Linux/Mac 下為Python函數(shù)添加超時時間
我們在使用 requests 這類網(wǎng)絡(luò)請求第三方庫時,可以看到它有一個參數(shù)叫做timeout,就是指在網(wǎng)絡(luò)請求發(fā)出開始計算,如果超過 timeout 還沒有收到返回,就拋出超時異常。(當(dāng)然存在特殊情況timeout 會失效,請看Timeouts and cancellation for humans*[1] 這篇文章中作者的舉例,我們不考慮這種特殊情況)。
但大家有沒有考慮過,如何為普通的函數(shù)設(shè)置超時時間?特別是在運行一些數(shù)據(jù)處理、AI 相關(guān)的代碼時,某個函數(shù)可能會運行很長時間,我們想實現(xiàn),在函數(shù)運行超過特定的時間時,自動報錯。
例如有這樣一個場景,我寫了一個函數(shù)calc_statistic(datas),根據(jù)用戶傳入的數(shù)據(jù)計算某個值。但如果用戶傳入的數(shù)據(jù)非常大,這個函數(shù)就可能運行很長時間。我想設(shè)置讓這個函數(shù)最多運行10秒鐘。如果10秒還沒有運行完成,就報錯。應(yīng)該怎么辦呢?
如果你的電腦操作系統(tǒng)是 Linux 或者 macOS,那么 可以使用 signal 來解決。
在公眾號前幾天的文章中,我們介紹了使用signal來接管鍵盤的中斷信號:《一日一技:在 Python 中接管鍵盤中斷信號》,用到的是signal.SIGINT。今天我們要用到的是signal.SIGALRM。
首先我們來看看這個信號的使用方法:
- import time
- import signal
- def handler(signum, _):
- print('定時到!')
- raise Exception('定時到了!')
- def clac_statistic(datas):
- time.sleep(100)
- signal.signal(signal.SIGALRM, handler)
- signal.alarm(5)
- clac_statistic('xxx')
運行效果如下圖所示:
首先綁定signal.SIGALRM事件到handler函數(shù)中,然后使用signal.alarm(10)延遲10秒發(fā)送一個信號。10秒到了以后,函數(shù)handler被運行。在函數(shù)中拋出了一個異常,導(dǎo)致程序結(jié)束。clac_statistic函數(shù)原本要運行100秒,但是在10秒以后就停止了,從而實現(xiàn)了函數(shù)的超時功能。
基于以上原理,我們實現(xiàn)一個裝飾器,來簡化為不同函數(shù)設(shè)置超時功能:
- import time
- import signal
- class FuncTimeoutException(Exception):
- pass
- def handler(signum, _):
- raise FuncTimeoutException('函數(shù)定時到了!')
- def func_timeout(times=0):
- def decorator(func):
- if not times:
- return func
- def wraps(*args, **kwargs):
- signal.alarm(times)
- result = func(*args, **kwargs)
- signal.alarm(0) # 函數(shù)提前運行完成,取消信號
- return result
- return wraps
- return decorator
- signal.signal(signal.SIGALRM, handler)
我們來試一試測試一下這個函數(shù)超時裝飾器。首先測試函數(shù)的運行時間小于超時時間時,程序正常運行沒有問題:
再來測試一下函數(shù)運行時間超過超時時間的情況:
正常拋出FuncTimeoutException異常。
那我們在實際使用中,可以使用try...except FuncTimeoutException捕獲這個異常,然后實現(xiàn)自定義的處理流程,例如:
- try:
- clac_statistic(100)
- except FuncTimeException:
- print('該函數(shù)運行超時,運行自定義的處理流程')
當(dāng)然你如果想直接跳過這個異常也沒問題,參考《一日一技:不使用 try...except 掩蓋一些已知異?!?/a>:
- import contextlib:
- with contextlib.supress(FuncTimeException):
- clac_statistic(100)