如何编写 Python 装饰器来增加堆栈深度？

Question

背景

在玩的时候，我经常写一些简单的递归函数，看起来像：

def f(a,b):
    if a>=0 and b>=0:
        return min( f(a-1,b) , f(b,a-1) ) # + some cost that depends on a,b
    else:
        return 0

（例如，在计算加权编辑距离或评估递归定义的数学公式时。）

然后我使用记忆装饰器自动缓存结果。

问题

当我尝试类似 f(200,10) 时，我得到：

RuntimeError: maximum recursion depth exceeded

这符合预期，因为递归实现耗尽了 Python 的堆栈空间/递归限制。

解决方法

我通常通过以下方法之一解决此问题：

• 使用 sys.setrecursionlimit 增加递归限制（仅适用于大约 1000 深度）
• 使用 for 循环为较小的值填充缓存
• 更改函数以将列表用作手动堆栈（通过 append 和 pop 调用）（换句话说，从递归实现转移到迭代实现）
• 使用其他编程语言

但我发现所有这些都很容易出错。

问题

有没有办法编写一个@Bigstack 装饰器来模拟拥有一个非常大的堆栈的效果？

请注意，我的函数通常会进行多次递归函数调用，因此这与尾递归不同 - 我确实希望将每个函数的所有内部状态保存在堆栈中。

我尝试了什么

我一直在考虑使用生成器表达式列表作为我的堆栈。通过探测堆栈帧，我可以确定函数何时被递归调用，然后触发异常以返回装饰器代码。但是，我无法想出一种方法将这些想法粘合在一起，以制造出真正可行的东西。

或者，我可以尝试访问函数的抽象语法树，并尝试将对递归函数的调用转换为 yield 语句，但这似乎走错了方向。

有什么建议吗？

编辑

看起来我确实在滥用 Python，但我一直在考虑的另一种方法是为每个块（例如 500 个堆栈帧）使用不同的线程，然后在每对连续的线程之间插入队列 - 一个队列用于参数，以及另一个返回值队列。 （每个队列最多有一个条目。）我认为这可能由于某种原因不起作用 - 但我可能只有在尝试实现它之后才能弄清楚为什么。

Answer 1

要绕过递归限制，您可以捕获RuntimeError 异常以检测堆栈空间何时用完，然后返回一个连续函数，当调用该函数时，将在以下级别重新启动递归你的空间用完了。调用它（及其返回值，等等），直到你得到一个值，然后从顶部重试。一旦你记住了较低的级别，较高的级别就不会遇到递归限制，所以最终这会起作用。将重复调用直到它工作在包装函数中。基本上，这是你预热缓存想法的一个懒惰版本。

这是一个简单的递归“从 1 到 n 加数字”函数的示例。

import functools

def memoize(func):
    cache = {}
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        key = args, tuple(sorted(kwargs.items()))
        if key in cache:
            return cache[key]
        else:
            result = func(*args, **kwargs)
            if not callable(result):
                cache[key] = result
            return result
    return wrapper

@memoize
def _addup(n):
    if n < 2:
        return n
    else:
        try:
            result = _addup(n - 1)
        except RuntimeError:
            return lambda: _addup(n)
        else:
            return result if callable(result) else result + n

def addup(n):
    result = _addup(n)
    while callable(result):
        while callable(result):
            result = result()
        result = _addup(n)
    return result

assert addup(5000) == sum(xrange(5001))

与其在调用链中一直返回 lambda 函数，我们可以引发异常以将其短路，这既提高了性能又简化了代码：

# memoize function as above, or you can probably use functools.lru_cache

class UnwindStack(Exception):
    pass

@memoize
def _addup(n):
    if n < 2:
        return n
    else:
        try:
            return _addup(n - 1) + n
        except RuntimeError:
            raise UnwindStack(lambda: _addup(n))

def _try(func, *args, **kwargs):
    try:
        return func(*args, **kwargs)
    except UnwindStack as e:
        return e[0]

def addup(n):
    result = _try(_addup, n)
    while callable(result):
        while callable(result):
            result = _try(result)
        result = _try(_addup, n)
    return result

不过，这仍然很不优雅，并且仍然有相当多的开销，我无法想象你会如何制作一个装饰器。我猜 Python 不太适合这种事情。

Answer 2

这是一个使用生成器表达式列表作为堆栈的实现：

def run_stackless(frame):
    stack, return_stack = [(False, frame)], []
    while stack:
        active, frame = stack.pop()
        action, res = frame.send(return_stack.pop() if active else None)
        if action == 'call':
            stack.extend([(True, frame), (False, res)])
        elif action == 'tail':
            stack.append((False, res))
        elif action == 'return':
            return_stack.append(res)
        else:
            raise ValueError('Unknown action', action)
    return return_stack.pop()

要使用它，您需要根据以下规则转换递归函数：

 return expr                    -> yield 'return', expr
 recursive_call(args...)        -> (yield 'call', recursive_call(args...))
 return recursive_call(args...) -> yield 'tail', recursive_call(args...)

例如，成本函数为a * b，您的函数变为：

def f(a,b):
    if a>=0 and b>=0:
        yield 'return', min((yield 'call', f(a-1,b)),
                            (yield 'call', f(b,a-1))) + (a * b)
    else:
        yield 'return', 0

测试：

In [140]: run_stackless(g(30, 4))
Out[140]: 410

在 Python 2.6.2 中，与直接调用相比，它似乎提供了约 8-10 倍的性能损失。

tail 动作用于尾递归：

def factorial(n):
    acc = [1]
    def fact(n):
        if n == 0:
            yield 'return', 0
        else:
            acc[0] *= n
            yield 'tail', fact(n - 1)
    run_stackless(fact(n))
    return acc[0]

向生成器递归样式的转换相当容易，并且可以通过字节码破解来完成。

Answer 3

这种方法将记忆和增加的堆栈深度结合到一个装饰器中。

我生成了一个线程池，每个线程负责 64 层堆栈。
线程只创建一次并重新使用（但目前从未删除）。

队列用于在线程之间传递信息，但请注意，只有当前堆栈深度对应的线程才会真正有工作要做。

我的实验表明，这为简单的递归函数增加了大约 10% 的开销（对于更复杂的函数应该更少）。

import threading,Queue

class BigstackThread(threading.Thread):

    def __init__(self,send,recv,func):
        threading.Thread.__init__( self )
        self.daemon = True
        self.send = send
        self.recv = recv
        self.func = func

    def run(self):
        while 1:
            args = self.send.get()
            v = self.func(*args)
            self.recv.put(v)

class Bigstack(object):

    def __init__(self,func):
        self.func = func
        self.cache = {}
        self.depth = 0
        self.threadpool = {}

    def __call__(self,*args):
        if args in self.cache:
            return self.cache[args]
        self.depth+=1
        if self.depth&63:
            v = self.func(*args)
        else:
            T=self.threadpool
            if self.depth not in T:
                send = Queue.Queue(1)
                recv = Queue.Queue(1)
                t = BigstackThread(send,recv,self)
                T[self.depth] = send,recv,t
                t.start()
            else:
                send,recv,_ = T[self.depth]
            send.put(args)
            v = recv.get()

        self.depth-=1
        self.cache[args]=v
        return v

@Bigstack
def f(a,b):
    if a>=0 and b>=0:
        return min(f(a-1,b),f(b-1,a))+1
    return 0