如何以比较优雅的方式为python-while增加超时？

遇到的问题：

在日常编码过程，可能会遇到这种场景：需要较多次尝试获取数据，你写的程序不知道何时能拿到数据，于是通过for或者while去“询问”，如以下的伪代码示意：

def get_data_from_remote():
      return data

while True:
      data = get_data_from_remote()
      if data:
          break

这段代码，只要没有获取到数据，就会一直循环执行，主线程一直处于阻塞状态，也就是我们常说的“死循环”。那么我们怎么解决这一问题呢？

几种方案：

如果作为初级玩家，每次遇到这种场景，都写以下一段逻辑：

start_time = datetime.datetime.now()
timeout = 5
while True:
      data = get_data_from_remote()
      if data:
          break
      if (datetime.datetime.now() -start_time).seconds> timeout:
            raise TimeoutException

不错效果达到了，但不具有通用性，当然小编想不到更低级的写法了。所以思考一下，怎么更具通用性？
如果 while关键字能自带超时效果就好了，以小编浅薄的知识面，想到了两种方式：

python解释器
修改python解释器，把while这个关键字功能丰富一下（虽然C语言是最好的语言，但我的这块知识都全部还给大学老师了）
ast
通过ast（抽象语法树）这个底层库，通过代码替换的方式，这种方式小编几年前用过，比如把time.sleep(1000) 悄悄地改成sleep(1)，再例如python著名测试框架pytest就是用这种方式把assert关键字给重写了。

然而这两种方式都要较深的功力，第2种方式可以做出来，但实现成本也不低，性能也不会太高。

方式3
于是思考有没有更简单的方式，想到可以“自动”修改while的条件, 请看下面的代码：

# self_while.py
import datetime
import sys
import traceback
from contextlib import contextmanager


class Condition(object):
    def __init__(self, cond_str, timeout):
        self.start = datetime.datetime.now()
        self.cond_str = cond_str
        self.timeout = timeout

    def __call__(self, *args, **kwargs):
        frame = sys._getframe(1)
        c = eval(self.cond_str, frame.f_locals)
        return (datetime.datetime.now() - self.start).seconds < self.timeout and c


@contextmanager
def ext_while_cond(whl_cond, timeout=5):
    cond_str = get_cond_str()
    cond = Condition(cond_str, timeout)
    yield cond


whl_line_ = 'with ' + ext_while_cond.__name__ + '('


def get_cond_str():
    x = traceback.extract_stack()
    for i in range(len(x)):
        code_line = x[i].line
        if whl_line_ in code_line:
            break
    if whl_line_ not in code_line:
        raise Exception('Cannot Find While Statement')
    else:
        l_idx, r_idx = code_line.find("("), code_line.rfind(")")
        return code_line[l_idx + 1, r_idx]

使用这些代码的例子：

import time

from self_while import ext_while_cond


def test1():
    flag = True
    with ext_while_cond(flag is True) as cond:
        while cond():
            # flag = False
            time.sleep(1)
            print("A")


def test2():
    flag1 = True
    flag2 = False
    with ext_while_cond((flag1 and not flag2), timeout=2) as cond:
        while cond():
            time.sleep(1)
            print("A")


if __name__ == '__main__':
    test1()

可以看出，调用的地方比之前就多一行代码：with ext_while_cond((flag1 and not flag2), timeout=2) as cond。
self_while里面的代码，使用了反射、作用域、上下文管理装饰器等知识，可以细品一下，全展开说太累了。

这个工具类没有实现超时后，抛一个超时异常出来，读者可以自己实现。
感谢阅读。

最后编辑于：2020.11.07 20:38:27