python - 在 python 中使用上下文管理器使进程超时

Question

我正在尝试编写一个上下文管理器来使进程超时，并且我有以下内容:

import threading
import _thread
import time
from contextlib import contextmanager

@contextmanager
def raise_timeout(timeout):
    timer = threading.Timer(timeout, _thread.interrupt_main)
    timer.start()

    try:
        yield
    except:
        print(f"timeout after {timeout} seconds")
        raise
    finally:
        timer.cancel()

with raise_timeout(1):
    time.sleep(5)

问题是当达到超时时，它仍然会等待 time.sleep(5) 完成，然后才会引发异常。当我运行脚本时，我等待了 5 秒钟，然后在终端中获得以下输出:

> timeout after 1 seconds Traceback (most recent call last):   File
> "scratch1.py",
> line xx, in <module>
>     time.sleep(5) KeyboardInterrupt

但是如果我像下面这样直接运行计时器，我会在 1 秒超时后立即收到异常。

timer = threading.Timer(1, _thread.interrupt_main)
timer.start()

我只是想知道 with 语句中的进程如何阻止异常直到它完成？

Answer 1

我认为您的 Timer 解决方案无法正常工作，因为 time.sleep 在调用 _thread.interrupt_main() 时不会被终止。

_thread文档对它的功能提出了建议:

Threads interact strangely with interrupts: the KeyboardInterrupt exception will be received by an arbitrary thread. (When the signal module is available, interrupts always go to the main thread.)

尽管 issue on Python's bug tracker建议 _thread.interrupt_main() 的文档不精确。

应该注意的是，在一个长时间运行的进程中替换 time.sleep 您的上下文管理器会按预期工作。

尝试:

with raise_timeout(1):
    for _ in range(1_000_000_000):
        pass

它会引发一个KeyboardInterrupt

---

如果便携性不是问题，这里有一个使用 signal 的解决方案仅在 Unix 上可用，因为它使用 SIGALARM。

import signal
import time
from contextlib import contextmanager

class TimeOutException(Exception):
    pass

@contextmanager
def raise_timeout(timeout):
    def _handler(signum, frame):
        raise TimeOutException()
    signal.signal(signal.SIGALRM, _handler)
    signal.alarm(timeout)

    try:
        yield
    except TimeOutException:
        print(f"Timeout after {timeout} seconds")
        raise
    finally:
        signal.alarm(0)

with raise_timeout(1):
    time.sleep(2)

这将引发 TimeOutException