pythonpython-3.xgeneratormetaprogramming

calculating time of generator of generic function


Trying to calculate the time that each yield takes from any function using generator, but I cant find any way to do so. The closest I got is creating a wrapper like this:

def catchtime_wrapper_for_method(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        return_val = func(*args, **kwargs)
        if isinstance(return_val, types.GeneratorType):
            for g in return_val:
                start_time = perf_counter()
                yield g
                end_time = perf_counter()
                print(f"Iteration time: {end_time - start_time:.6f} seconds")
        else:
            return return_val
    return wrapper

and calling like that:

for i in catchtime_wrapper_for_method(my_gen_function)(args...):
   # doing something with i here

my_gen_function yielding result of some io + calculation on the io. it has tqdm which is proven to be accurate and prints ~5sec per iteration.

But catchtime_wrapper_for_method prints ~1sec, I'm assuming ~1sec is the time to run catchtime_wrapper_for_method without the my_gen_function part... so i guess 6sec should be the accurate result but i need it to be printed by catchtime_wrapper_for_method and not the tqdm. I need my_gen_function continue to behave as always so it needs to still be a generator.

Also tried: (you are welcome to try too, I could have done some mistakes)

Any suggestion why it is happening?

Thanks :)


Solution

  • Seems like the difference arises cause tqdm is measuring the time for each iteration (including any overhead or processing within the generator function itself) whereas the wrapper function only measures the time spent in the wrapper function, not the generator function.

    Try something like this:

    def catchtime_wrapper_for_method(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            generator = func(*args, **kwargs)
            while True:
                try:
                    start_time = perf_counter()
                    result = next(generator)
                    end_time = perf_counter()
                    print(f"Iteration time: {end_time - start_time:.6f} seconds")
                    yield result
                except StopIteration:
                    break
    
        return wrapper
    

    For example, here's how you would call it:

    def my_gen_function(*args, **kwargs):
        for i in range(args[0]):
            print(args, kwargs)
            yield i
            
    for i in catchtime_wrapper_for_method(my_gen_function)(23, 45, something="something"):
        print(i)