1. 程式人生 > >python裝飾器限制函式執行時間,超時退出

python裝飾器限制函式執行時間,超時退出

實際專案中會涉及到需要對有些函式的響應時間做一些限制,如果超時就退出函式的執行,停止等待。

可以利用python中的裝飾器實現對函式執行時間的控制。

python裝飾器簡單來說可以在不改變某個函式內部實現和原來呼叫方式的前提下對該函式增加一些附件的功能,提供了對該函式功能的擴充套件。

方法一. 使用 signal

# coding=utf-8
import signal
import time

def set_timeout(num, callback):
    def wrap(func):
        def handle(signum, frame):  # 收到訊號 SIGALRM 後的回撥函式,第一個引數是訊號的數字,第二個引數是the interrupted stack frame.
            raise RuntimeError

        def to_do(*args, **kwargs):
            try:
                signal.signal(signal.SIGALRM, handle)  # 設定訊號和回撥函式
                signal.alarm(num)  # 設定 num 秒的鬧鐘
                print('start alarm signal.')
                r = func(*args, **kwargs)
                print('close alarm signal.')
                signal.alarm(0)  # 關閉鬧鐘
                return r
            except RuntimeError as e:
                callback()

        return to_do

    return wrap

def after_timeout():  # 超時後的處理函式
    print("Time out!")

@set_timeout(2, after_timeout)  # 限時 2 秒超時
def connect():  # 要執行的函式
    time.sleep(3)  # 函式執行時間,寫大於2的值,可測試超時
    print('Finished without timeout.')

if __name__ == '__main__':
    connect()

方法一中使用的signal有所限制,需要在linux系統上,並且需要在主執行緒中使用。方法二使用執行緒計時,不受此限制。

方法二. 使用Thread

# -*- coding: utf-8 -*-
from threading import Thread
import time

class TimeoutException(Exception):
    pass

ThreadStop = Thread._Thread__stop

def timelimited(timeout):
    def decorator(function):
        def decorator2(*args,**kwargs):
            class TimeLimited(Thread):
                def __init__(self,_error= None,):
                    Thread.__init__(self)
                    self._error =  _error

                def run(self):
                    try:
                        self.result = function(*args,**kwargs)
                    except Exception,e:
                        self._error = str(e)

                def _stop(self):
                    if self.isAlive():
                        ThreadStop(self)

            t = TimeLimited()
            t.start()
            t.join(timeout)

            if isinstance(t._error,TimeoutException):
                t._stop()
                raise TimeoutException('timeout for %s' % (repr(function)))

            if t.isAlive():
                t._stop()
                raise TimeoutException('timeout for %s' % (repr(function)))

            if t._error is None:
                return t.result

        return decorator2
    return decorator

@timelimited(2)  # 設定執行超時時間2S
def fn_1(secs):
    time.sleep(secs)
    return 'Finished without timeout'

def do_something_after_timeout():
    print('Time out!')

if __name__ == "__main__":
    try:
        print(fn_1(3))  # 設定函式執行3S
    except TimeoutException as e:
        print(str(e))
        do_something_after_timeout()