Can you write a python decorator that works for generator functions and normal functions, and matches their return styles?

Question

Let's say I want to write a python decorator to time a function, and let the user pass in how many times they'd like it to run. I want this decorator to return on a function that returns, and I want it to return a generator if the decorated function uses yield statements.

If I do the following:

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

You will see that any function you decorate now returns a generator.

@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

Now, I can get it to match the return type by putting the yield into another helper function, like so:

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, *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

Now, each function has the expected return type:

@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

The problem now, is that of course two didn't really run. It's a generator, so I didn't get a meaningful time read off it!

Does anyone know of a way to have a decorator match the return type (normal return vs generator), and also meaningfully measure time?

To be clear, this is an example problem to demonstrate the questions I have about decorators and python. I know there are many open source tools for timing things, it's more of an illustrative problem.

Answer 1

As you noticed, part of the problem with trying to combine them is as soon as a yield occurs within a function, the function becomes a generator. You can get around this by breaking it down into two separate functions, where one handles "normal" functions, and the other handles generators. Your attempt of trying to share the timing code between the two implementations would be very difficult to get right though because each case needs to be handled differently.

In the generator wrapper, I'm using a bare for loop (for _ in func(*args, **kwargs): pass) to "force" the generator so that it can be properly timed. Note how I'm throwing away the results that it returns. After iters-1 iterations though, I make a separate call to func, and this time I actually use the results.

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

And an example of the use:

@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

Gives:

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'], ['a', 'b'], ['a', 'b'], ['a', 'b'], ['a', 'b'], ['a', 'b'], ['a', 'b'], ['a', 'b'], ['a', 'b']]

There's an unfortunate amount of duplication in the wrappers, but as I noted above, that's difficult to avoid, and isn't a huge problem.

---

Note though that this will only time your generator function properly if you manually force the returned generator (like by using list), so it may not make sense in the end to return a generator if you intend on timing the function. Unless you want to time how long the generator takes to complete after being run for the first time (for curiosity's sake maybe), you should probably evaluate the generator before returning to ensure accurate results.