我們?cè)谑褂?requests 這類網(wǎng)絡(luò)請(qǐng)求第三方庫(kù)時(shí)，可以看到它有一個(gè)參數(shù)叫做 timeout ，就是指在網(wǎng)絡(luò)請(qǐng)求發(fā)出開始計(jì)算，如果超過 timeout 還沒有收到返回，就拋出超時(shí)異常。(當(dāng)然存在特殊情況timeout 會(huì)失效，請(qǐng)看Timeouts and cancellation for humans* 這篇文章中作者的舉例，我們不考慮這種特殊情況）。

但大家有沒有考慮過，如何為普通的函數(shù)設(shè)置超時(shí)時(shí)間？特別是在運(yùn)行一些數(shù)據(jù)處理、AI 相關(guān)的代碼時(shí)，某個(gè)函數(shù)可能會(huì)運(yùn)行很長(zhǎng)時(shí)間，我們想實(shí)現(xiàn)，在函數(shù)運(yùn)行超過特定的時(shí)間時(shí)，自動(dòng)報(bào)錯(cuò)。

例如有這樣一個(gè)場(chǎng)景，我寫了一個(gè)函數(shù) calc_statistic(datas) ，根據(jù)用戶傳入的數(shù)據(jù)計(jì)算某個(gè)值。但如果用戶傳入的數(shù)據(jù)非常大，這個(gè)函數(shù)就可能運(yùn)行很長(zhǎng)時(shí)間。我想設(shè)置讓這個(gè)函數(shù)最多運(yùn)行10秒鐘。如果10秒還沒有運(yùn)行完成，就報(bào)錯(cuò)。應(yīng)該怎么辦呢？

如果你的電腦操作系統(tǒng)是 Linux 或者 macOS，那么 可以使用 signal 來解決。

在公眾號(hào)前幾天的文章中，我們介紹了使用signal來接管鍵盤的中斷信號(hào)，用到的是 signal.SIGINT 。今天我們要用到的是 signal.SIGALRM 。

首先我們來看看這個(gè)信號(hào)的使用方法：

import time
import signal
def handler(signum, _):
  print('定時(shí)到！')
  raise Exception('定時(shí)到了！')
def clac_statistic(datas):
  time.sleep(100)
signal.signal(signal.SIGALRM, handler)
signal.alarm(5)
clac_statistic('xxx')

運(yùn)行效果如下圖所示：

首先綁定 signal.SIGALRM 事件到 handler 函數(shù)中，然后使用 signal.alarm(10) 延遲10秒發(fā)送一個(gè)信號(hào)。10秒到了以后，函數(shù) handler 被運(yùn)行。在函數(shù)中拋出了一個(gè)異常，導(dǎo)致程序結(jié)束。 clac_statistic 函數(shù)原本要運(yùn)行100秒，但是在10秒以后就停止了，從而實(shí)現(xiàn)了函數(shù)的超時(shí)功能。

基于以上原理，我們實(shí)現(xiàn)一個(gè)裝飾器，來簡(jiǎn)化為不同函數(shù)設(shè)置超時(shí)功能：

import time
import signal
class FuncTimeoutException(Exception):
  pass
def handler(signum, _):
  raise FuncTimeoutException('函數(shù)定時(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ù)提前運(yùn)行完成，取消信號(hào)
      return result
    return wraps
  return decorator
signal.signal(signal.SIGALRM, handler)

我們來試一試測(cè)試一下這個(gè)函數(shù)超時(shí)裝飾器。首先測(cè)試函數(shù)的運(yùn)行時(shí)間小于超時(shí)時(shí)間時(shí)，程序正常運(yùn)行沒有問題：

再來測(cè)試一下函數(shù)運(yùn)行時(shí)間超過超時(shí)時(shí)間的情況：

正常拋出 FuncTimeoutException 異常。

那我們?cè)趯?shí)際使用中，可以使用 try...except FuncTimeoutException 捕獲這個(gè)異常，然后實(shí)現(xiàn)自定義的處理流程，例如：

try:
  clac_statistic(100)
except FuncTimeException:
  print('該函數(shù)運(yùn)行超時(shí)，運(yùn)行自定義的處理流程')

當(dāng)然你如果想直接跳過這個(gè)異常也沒問題：

import contextlib:
with contextlib.supress(FuncTimeException):
  clac_statistic(100)

總結(jié)

以上所述是小編給大家介紹的在 Linux/Mac 下為Python函數(shù)添加超時(shí)時(shí)間的方法，希望對(duì)大家有所幫助，也非常感謝大家對(duì)腳本之家網(wǎng)站的支持！

最新評(píng)論