pythonの処理に制限時間を設定する

探索アルゴリズムを実装した際に、「10秒で探索！」など制限時間をつけて探索したいなと思うことがあり、色々探しているとtimeout-decoratorという良さげなライブラリがありました。

timeout-decorator

pipインストールすれば使えます。

制限時間が来ると強制終了するような使い方は参考になる記事があったのですが、探索アルゴリズムで使う際には制限時間がきたら最新の結果を返すような挙動をさせたいです。

ということで以下のようなコードになりました。

test

import time
import timeout_decorator


def long_function(i):
    time.sleep(4)
    return i


class time_limit_class:
    def __init__(self):
        self.result = 0

    @timeout_decorator.timeout(5)
    def main(self):
        for i in range(1000):
            self.result = long_function(i)
        return self.result

    def timeout(self):
        return self.result


if __name__ == "__main__":
    time_limit = time_limit_class()
    try:
        print(time_limit.main())
    except:
        print(time_limit.timeout())

やってることとしては、とりあえずインスタンス作ってそこに値を保持しておき、timeoutの際にはそれを返すといったものになっています。

正直雑な実装だなとは思っているので、もっと良い方法があれば教えてください〜。