在 Linux/Mac 下爲Python函數添加超時時間

咱們在使用 requests 這類網絡請求第三方庫時，能夠看到它有一個參數叫作timeout，就是指在網絡請求發出開始計算，若是超過 timeout 尚未收到返回，就拋出超時異常。(固然存在特殊狀況timeout 會失效，請看Timeouts and cancellation for humans* 這篇文章中做者的舉例，咱們不考慮這種特殊狀況）。

但你們有沒有考慮過，如何爲普通的函數設置超時時間？特別是在運行一些數據處理、AI 相關的代碼時，某個函數可能會運行很長時間，咱們想實現，在函數運行超過特定的時間時，自動報錯。

例若有這樣一個場景，我寫了一個函數calc_statistic(datas)，根據用戶傳入的數據計算某個值。但若是用戶傳入的數據很是大，這個函數就可能運行很長時間。我想設置讓這個函數最多運行10秒鐘。若是10秒尚未運行完成，就報錯。應該怎麼辦呢？

若是你的電腦操做系統是 Linux 或者 macOS，那麼 可使用 signal 來解決。

在公衆號前幾天的文章中，咱們介紹了使用signal來接管鍵盤的中斷信號，用到的是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函數中，而後使用signal.alarm(10)延遲10秒發送一個信號。10秒到了之後，函數handler被運行。在函數中拋出了一個異常，致使程序結束。clac_statistic函數本來要運行100秒，可是在10秒之後就中止了，從而實現了函數的超時功能。

基於以上原理，咱們實現一個裝飾器，來簡化爲不一樣函數設置超時功能：

import time
import signal


class FuncTimeoutException(Exception):
    pass

def handler(signum, _):
    raise FuncTimeoutException('函數定時到了！')

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)  # 函數提早運行完成，取消信號
            return result
        return wraps
    return decorator

signal.signal(signal.SIGALRM, handler)
複製代碼

咱們來試一試測試一下這個函數超時裝飾器。首先測試函數的運行時間小於超時時間時，程序正常運行沒有問題：

再來測試一下函數運行時間超過超時時間的狀況：

正常拋出FuncTimeoutException異常。

那咱們在實際使用中，可使用try...except FuncTimeoutException捕獲這個異常，而後實現自定義的處理流程，例如：

try:
    clac_statistic(100)
except FuncTimeException:
    print('該函數運行超時，運行自定義的處理流程')
複製代碼

固然你若是想直接跳過這個異常也沒問題：

import contextlib:
with contextlib.supress(FuncTimeException):
    clac_statistic(100)
複製代碼