装饰器是否有办法将下面的函数转换为生成器?
@decorator_that_makes_func_into_generator
def countdown(n):
while n > 0:
print n,
n = n - 1
如有必要,可以修改该功能。请注意,该函数没有 yield 语句,否则它已经是一个生成器。
请您参考如下方法:
如果您无法更改countdown 的源,则必须捕获输出:
import sys
from io import StringIO
def decorator_that_makes_func_into_generator(func):
def wrapper(*a, **ka):
# Temporarily redirect all output to StringIO instance (intr)
ts, intr = sys.stdout, StringIO()
sys.stdout = intr
func(*a, **ka)
# Restore normal stdout from backup (ts)
sys.stdout = ts
# Get output from intr, split it by whitespace and use it as generator
yield from intr.getvalue().split()
return wrapper
@decorator_that_makes_func_into_generator
def countdown(n):
while n > 0:
print(n)
n = n - 1
print(countdown(5), list(countdown(5)))
# <generator object wrapper at 0x01E09058> ['5', '4', '3', '2', '1']
如果您可以更改函数,并且如果您想从countdown(list 或其他序列类型)返回一些内容,并且然后从返回的对象创建一个生成器,装饰器看起来像
def decorator_that_makes_func_into_generator(func):
def wrapper(*a, **ka):
yield from func(*a, **ka)
return wrapper
注意:很棒的 yield from 是在 Python 3.3 中引入的,在旧版本上使用普通循环代替:
for x in func(*a, **ka):
yield x
例子:
@decorator_that_makes_func_into_generator
def countdown(n):
res = []
while n > 0:
res.append(n)
n = n - 1
return res
print(type(countdown(5)), list(countdown(5)))
# Output: <class 'generator'> [5, 4, 3, 2, 1]
没有什么能阻止您将 decorator_that_makes_func_into_generator 应用于生成器:
@decorator_that_makes_func_into_generator
def countdown(n):
while n > 0:
yield n
n = n - 1
print(type(countdown(5)), list(countdown(5)))
# Outputs <class 'generator'> [5, 4, 3, 2, 1]
