Skip to main content
 首页 » 编程设计

python之在 Python 中使用装饰器将函数变成生成器

2025年05月04日14pengyingh

装饰器是否有办法将下面的函数转换为生成器?

@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]