使用装饰器将函数转换为Python中的生成器

时间:2014-05-16 22:03:43

标签: python generator decorator

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

@decorator_that_makes_func_into_generator
def countdown(n):
    while n > 0:
        print n,
        n = n - 1

如有必要,可以修改该功能。请注意,该函数没有yield语句,否则它已经是一个生成器。

2 个答案:

答案 0 :(得分:8)

如果您无法更改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']

如果您可以更改功能,并且想要从countdownlist或其他序列类型)和返回某些内容,那么创建来自返回对象的生成器,装饰器看起来像

def decorator_that_makes_func_into_generator(func):
    def wrapper(*a, **ka):
        yield from func(*a, **ka)
    return wrapper

注意:在Python 3.3中引入了令人敬畏的yield from,旧版本使用了普通循环:

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]

答案 1 :(得分:5)

如果你不能改变那个功能的主体,我担心这太难了。

您尝试换行的函数不是生成器,即使将其包装到生成器中,此函数也将从头开始一次执行。大概你不想要那个。

理论上你能做什么?

  • 在调试器下运行
  • 逐行,逐行运行
  • 访问函数源,修改它,编译为字节码
  • 修改字节码
  • 覆盖print()(在Python3中更容易)
  • 使用stackless python在任意点保存和恢复堆栈