python - 你能写一个适用于生成器函数和普通函数并匹配它们的返回样式的 python 装饰器吗？

Question

假设我想编写一个 python 装饰器来为函数计时，并让用户输入他们希望它运行的次数。我希望这个装饰器在返回的函数上返回，如果装饰函数使用 yield 语句，我希望它返回一个生成器。

如果我执行以下操作:

import time
from datetime import datetime
import inspect


def time_it(iters=1):
    def decorator(func):
        def wrapper(*args, **kwargs):
            is_gen = inspect.isgeneratorfunction(func)

            start = datetime.now()
            for _ in range(iters):
                ret = yield from func(*args, **kwargs) if is_gen else func(*args, **kwargs)

            elapsed = datetime.now() - start
            print(f'Elapsed time: {elapsed} over {iters} iterations')

            return ret
        return wrapper
    return decorator

您会看到您装饰的任何函数现在都会返回一个生成器。

@time_it()
def one(ret):
    time.sleep(1)
    return ret


@time_it()
def two(ret):
    time.sleep(1)
    yield from ret

x = one(['a', 'b'])
y = two(['a', 'b'])
print(type(x))  # generator

现在，我可以通过将 yield 放入另一个辅助函数来让它匹配返回类型，如下所示:

import time
from datetime import datetime
import inspect

def do_gen(func, *args, **kwargs):
    yield from func(*args, **kwargs)


def time_it(iters=1):
    def decorator(func):
        def wrapper(*args, **kwargs):
            is_gen = inspect.isgeneratorfunction(func)

            start = datetime.now()
            for _ in range(iters):
                ret = do_gen(func, *args, **kwargs) if is_gen else func(*args, **kwargs)

            elapsed = datetime.now() - start
            print(f'Elapsed time: {elapsed} over {iters} iterations')

            return ret
        return wrapper
    return decorator

现在，每个函数都有预期的返回类型:

@time_it()
def one(ret):
    time.sleep(1)
    return ret


@time_it()
def two(ret):
    time.sleep(1)
    yield from ret


one = one(['a', 'b'])
two = two(['a', 'b'])

print(type(one))  # list
print(type(two))  # generator

现在的问题是，two 当然没有真正运行。这是一个生成器，所以我没有花有意义的时间阅读它!

有谁知道让装饰器匹配返回类型(正常返回与生成器)并且有意义地测量时间的方法吗？

需要说明的是，这是一个示例问题，用于演示我对装饰器和 Python 的疑问。我知道有许多用于计时的开源工具，它更像是一个说明性问题。

Answer 1

正如您所注意到的，尝试组合它们的部分问题在于，一旦 yield 发生在函数中，该函数就会变成生成器。您可以通过将其分解为两个单独的函数来解决这个问题，其中一个处理“正常”函数，另一个处理生成器。但是，您尝试在两种实现之间共享计时代码的尝试很难做到正确，因为每种情况都需要以不同的方式处理。

在生成器包装器中，我使用了一个简单的 for 循环(for _ in func(*args, **kwargs): pass)来“强制”发电机，以便它可以正确计时。请注意我是如何丢弃它返回的结果的。不过，在 iters-1 迭代之后，我单独调用了 func，这次我实际使用了结果。

def time_it(iters=1):
    def decorator(func):
        def gen_wrapper(*args, **kwargs):
            start = datetime.now()
            for _ in range(iters - 1):
                for _ in func(*args, **kwargs):  # Force the generator for first the iters-1 tests
                    pass
            yield from func(*args, **kwargs)  # Then actually use the last result
            elapsed = datetime.now() - start
            print(f'Gen Elapsed time: {elapsed} over {iters} iterations')

        def func_wrapper(*args, **kwargs):
            start = datetime.now()
            ret = None
            for _ in range(iters):
                ret = func(*args, **kwargs)
            elapsed = datetime.now() - start
            print(f'Normal Elapsed time: {elapsed} over {iters} iterations')
            return ret
        return gen_wrapper if inspect.isgeneratorfunction(func) else func_wrapper
    return decorator

以及一个使用示例:

@time_it(2)
def one(ret):
    time.sleep(1)
    return ret


@time_it(2)
def two(ret):
    for n in range(10):
        yield ret
        time.sleep(0.5)

result = one(['a', 'b'])
print(type(result))
print(result)

gen = two(['a', 'b'])
print(type(gen))
print(list(gen))  # Using list to force the generator

给予:

Normal Elapsed time: 0:00:02.002115 over 2 iterations
<class 'list'>
['a', 'b']
<class 'generator'>
Gen Elapsed time: 0:00:10.037192 over 2 iterations
[['a', 'b'], ['a', 'b'], ['a', 'b'], ['a', 'b'], ['a', 'b'], ['a', 'b'], ['a', 'b'], ['a', 'b'], ['a', 'b'], ['a', 'b']]

不幸的是，包装器中存在大量重复，但正如我上面提到的，这是很难避免的，也不是什么大问题。

---

请注意，如果您手动强制返回生成器(例如使用 list)，这只会正确地计时您的生成器功能，因此如果您最终返回生成器可能没有意义打算计时的功能。除非您想计算生成器在第一次运行后需要多长时间才能完成(可能出于好奇)，否则您应该在返回之前评估生成器以确保准确的结果。