Python 装饰器用于计时递归函数

Question

我有一个简单的装饰器，用于跟踪函数调用的运行时间：

def timed(f):
    def caller(*args):
        start = time.time()
        res = f(*args)
        end = time.time()
        return res, end - start
    return caller

这可以如下所示使用，并返回函数结果和执行时间的元组。

@timed
def test(n):
    for _ in range(n):
        pass
    return 0

print(test(900)) # prints (0, 2.69e-05)

很简单。但是现在我想将其应用于递归函数。将上述包装器应用于递归函数会产生嵌套的元组，其中包含每个递归调用的次数，这是预期的结果。

@timed
def rec(n):
    if n:
        return rec(n - 1)
    else:
        return 0

print(rec(3)) # Prints ((((0, 1.90e-06), 8.10e-06), 1.28e-05), 1.90e-05)

有什么优雅的方法可以编写装饰器，以便它正确处理递归？显然，如果是一个定时函数，你可以包装调用：

@timed
def wrapper():
    return rec(3)

这将返回一个结果和时间的元组，但我希望装饰器来处理所有内容，这样调用者就不必担心为每个调用定义一个新函数。有什么好的想法吗？

- Channing Hurley

---

您可能可以使用sys._getframe来实现您想要的功能，但这是CPython的一个实现细节，因此它不是Python语言的真正一部分。此外，它可能会很慢。 - Bakuriu

Answer 1

这里的问题并不是装饰器本身，而是 rec 需要成为一种行为方式的函数，但你想让 rec 成为另一种行为方式的函数。没有简单的方法来将这两种行为与单个的 rec 函数协调一致。

最干净的选项是停止要求 rec 同时具有两种功能。不要使用装饰器符号，而是将 timed(rec) 分配给另一个名称：

def rec(n):
    ...

timed_rec = timed(rec)

如果您不想使用两个名称，则需要编写rec以理解经过装饰的rec将返回的实际值。例如：

@timed
def rec(n):
    if n:
        val, runtime = rec(n-1)
        return val
    else:
        return 0

- user2357112

---

是的，我想既然装饰器的目的是在每次调用函数时被调用，它本质上并不适用于递归。 - Channing Hurley

Answer 2

我更喜欢目前为止的其他答案（特别是user2357112的答案），但您也可以创建一个基于类的装饰器，检测函数是否已被激活，如果是，则绕过计时：

import time

class fancy_timed(object):
    def __init__(self, f):
        self.f = f
        self.active = False

    def __call__(self, *args):
        if self.active:
            return self.f(*args)
        start = time.time()
        self.active = True
        res = self.f(*args)
        end = time.time()
        self.active = False
        return res, end - start


@fancy_timed
def rec(n):
    if n:
        time.sleep(0.01)
        return rec(n - 1)
    else:
        return 0
print(rec(3))

(class written with (object) so that this is compatible with py2k and py3k).

需要注意的是，为了使其正常运行，最外层调用应该使用try和finally。这里是__call__的优化版本:

def __call__(self, *args):
    if self.active:
        return self.f(*args)
    try:
        start = time.time()
        self.active = True
        res = self.f(*args)
        end = time.time()
        return res, end - start
    finally:
        self.active = False

- torek

---

这个程序存在线程安全和可重入性问题。例如，在一个线程中调用rec可能会导致另一个线程中的调用省略时间信息。 - user2357112

@user2357112：同意；这就是我为什么更喜欢你的回答的原因。 :-) （可以获取当前线程信息，从而与threading模块配合使用，但是会遇到绿色线程等问题。） - torek

Answer 3

你可以通过*咳嗽*滥用contextmanager和function attribute，以不同的方式构建你的计时器...

from contextlib import contextmanager
import time

@contextmanager
def timed(func):
    timed.start = time.time()
    try:
        yield func
    finally:
        timed.duration = time.time() - timed.start

def test(n):
    for _ in range(n):
        pass
    return n

def rec(n):
    if n:
        time.sleep(0.05) # extra delay to notice the difference
        return rec(n - 1)
    else:
        return n

with timed(rec) as r:
    print(t(10))
    print(t(20))

print(timed.duration)

with timed(test) as t:
    print(t(555555))
    print(t(666666))

print(timed.duration)

结果：

# recursive
0
0
1.5130000114440918

# non-recursive
555555
666666
0.053999900817871094

如果被认为是不好的黑客行为，我很乐意接受你的批评。

- r.ook

---

滥用timed函数对象的属性来存储时间信息会导致问题，如果您尝试同时计时多个事物。此外，timed实际上不需要func，也没有真正计时func的执行；它正在计时with块的内容。 - user2357112

简而言之，我的回答不好吗？我最近才开始学习这些工具集，所以我认为这可能会起作用，但是我意识到自己有什么遗漏。我当时并没有考虑到多个计时器，所以你说得对。 - r.ook

Answer 4

当我尝试对一个简单的quicksort实现进行分析时，我遇到了同样的问题。

主要问题在于装饰器在每次函数调用时都会执行，我们需要一些可以保持状态的东西，以便在最后可以将所有调用加起来。装饰器不是完成这项工作的正确工具。

然而，一个想法是滥用函数是对象并且可以有属性的事实。下面通过一个简单的装饰器来探讨这个想法。必须理解的一点是，通过使用装饰器的语法糖（@），函数将始终累积其时间。

from typing import Any, Callable
from time import perf_counter

class timeit:

    def __init__(self, func: Callable) -> None:
        self.func = func
        self.timed = []

    def __call__(self, *args: Any, **kwds: Any) -> Any:
        start = perf_counter()
        res = self.func(*args, **kwds)
        end = perf_counter()        
        self.timed.append(end - start)
        return res

# usage

@timeit
def rec(n):
    ...

if __name__ == "__main__":
    result = rec(4) # rec result
    print(f"Took {rec.timed:.2f} seconds")
    # Out: Took 3.39 seconds
    result = rec(4) # rec result
    # timings between calls are accumulated
    # Out: Took 6.78 seconds

这让我们想到了一个受@r.ook启发的解决方案，下面是一个简单的上下文管理器，它存储每次运行时间并在结束时打印其总和（__exit__）。请注意，由于对于每个计时，我们都需要一个with语句，因此这不会累积不同的运行。

from typing import Any, Callable
from time import perf_counter

class timeit:

    def __init__(self, func: Callable) -> None:
        self.func = func
        self.timed = []

    def __call__(self, *args: Any, **kwds: Any) -> Any:
        start = perf_counter()
        res = self.func(*args, **kwds)
        end = perf_counter()
        self.timed.append(end - start)
        return res

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, exc_traceback):
        # TODO: report `exc_*` if an exception get raised
        print(f"Took {sum(self.timed):.2f} seconds")
        return

# usage

def rec(n):
    ...

if __name__ == "__main__":

    with timeit(rec) as f:
        result = f(a) # rec result
    # Out: Took 3.39 seconds

- Felipe Whitaker

Answer 5

虽然这并不是将递归与装饰器整合的总体解决方案，但对于仅涉及计时的问题，我已经验证了时间元组的最后一个元素是总运行时间，因为这是来自最上层递归调用的时间。因此，如果您有

@timed
def rec():
    ...

如果要获取原始函数定义的总运行时间，你可以简单地执行以下操作

rec()[1]

另一方面，获取调用结果则需要通过嵌套元组进行递归。

def get(tup):
    if isinstance(tup, tuple):
        return get(tup[0])
    else:
        return tup

这可能太复杂了，仅仅为了得到您的函数结果。

- Channing Hurley