指定時間内に関数が終了しなかったら何かするデコレータ

昨日投稿したTipsの改良版です。通知する部分をハンドラ関数を渡すようにして分離し任意の処理を行えるようにしました。

また、@methaneさんに教えていただいた functools.wrap を使ってみました。このデコレータがついた関数をpaverの@taskにもしてみたのですがヘルプがNoneとなり docstring が表示されなかったのはこれを使っていなかったからのようでした！(functools.wrapを使うとヘルプにdocstringが表示されました)

timeout.py

from functools import wraps

def on_timeout(limit, handler, hint=None):
    '''                                                                                                    
    指定した実行時間に終了しなかった場合、handlerをhint/limitを引数にして呼び出します                
    @on_timeout(limit=3600, handler=notify_func, hint=u'長い計算')                                         
    def long_time_function():                                                                              
    '''
    def notify_handler(signum, frame):
        handler("'%s' is not finished in %d second(s)." % (hint, limit))

    def __decorator(function):
        def __wrapper(*args, **kwargs):
            import signal
            signal.signal(signal.SIGALRM, notify_handler)
            signal.alarm(limit)
            result = function(*args, **kwargs)
            signal.alarm(0)
            return result
        return wraps(function)(__wrapper)
    return __decorator

実際の使い方は以下のような感じです

main.py

from timeout import on_timeout

def handler_func(msg):
    print msg # 実際は適切な通知処理を行う

@on_timeout(limit=1, handler=handler_func, hint=u'長い計算')
def long_time_function():
    import time
    time.sleep (2)

if __name__ == "__main__":
    long_time_function()
 

これ実行すると limit が1秒にも関わらず処理が sleep(2) しているため1秒経過時点でhandler_funcが呼ばれ、以下のような表示になります。

'長い計算' is not finished in 1 second(s).

これで AWS 以外でも応用できそうです。