user2979872 user2979872 - 1 month ago 5
Python Question

Use a decorator to turn a function into a generator in Python

Is there a way for a decorator to convert the function below into a generator?

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


The function can be modified if necessary. Note that the function does not have a yield statement, otherwise it would already be a generator.

Answer

If you can't change the source of countdown, you'll have to capture the output:

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

If you can change the function and if you want to return something from countdown (list or other sequence type) and then create a generator from the returned object, decorator'd look like

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

Note: the awesome yield from was introduced in Python 3.3, on older versions use plain loop instead:

for x in func(*a, **ka): 
    yield x

Example:

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

Nothing stops you from applying decorator_that_makes_func_into_generator to a 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]
Comments