您可以编写适用于生成器函数和普通函数并匹配其返回样式的python装饰器吗？

如何解决您可以编写适用于生成器函数和普通函数并匹配其返回样式的python装饰器吗？

假设我想编写一个python装饰器给一个函数计时，然后让用户传递他们希望它运行多少次。我希望这个装饰器在返回的函数上return，并且如果装饰的函数使用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,**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的疑问。我知道有很多用于计时的开源工具，这更多是一个说明性问题。

解决方法

您已经注意到，尝试将它们组合在一起的部分问题是，一旦在函数中出现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']]

不幸的是，包装程序中存在大量重复项，但是正如我上面指出的那样，这是很难避免的，而且不是一个大问题。

---

请注意，尽管只有手动强制返回的生成器（例如使用list），这才能使生成器功能正确计时，因此，如果您打算计时，最终返回生成器可能没有意义。功能。除非您想计时发生器首次运行后需要多长时间才能完成（也许是出于好奇的缘故），否则您应该在返回之前对发生器进行评估，以确保结果准确。