Unexpected decorator/generator behavior

Question

I created a timer decorator but it doesn't work if the decorated function is a generator.

import numpy as np
from time import time
from collections import Counter


def timer(f):
    def inner(*args, **kwargs):
        start = time()
        try:
            res = f(*args, **kwargs)
        except Exception as e:
            end = time()
            timer.counter.update({f'{f.__module__}.{f.__name__}': end - start})
            raise e
        end = time()
        timer.counter.update({f'{f.__module__}.{f.__name__}': end - start})
        return res
    return inner
timer.counter = Counter()


class AA:
    @timer
    def __init__(self):
        a = np.array(range(1_000_000))

    @timer
    def __iter__(self):
        a = np.array(range(1_000_000))
        yield 'a'

    @timer
    def normal_fun(self):
        a = np.array(range(1_000_000))

    @timer
    def fun_with_yield(self):
        a = np.array(range(1_000_000))
        yield 'a'


a = AA()
for i in a:
    pass
a.normal_fun()
a.fun_with_yield()
print(timer.counter)

Output:

Counter({'main.init': 0.10380005836486816, 'main.normal_fun': 0.10372400283813477, 'main.iter': 0.0, 'main.fun_with_yield': 0.0})

Why is the generator functions' time equal 0.0 and how can I fix it?

Answer 1

In the end I achieved expected results by creating another decorator for generators. The decorator is a generator in itself and counts time spent in each iteration.

def generator_timer(f):
    def inner(*args, **kwargs):
        start = time()
        try:
            res = f(*args, **kwargs)
            for i in res:
                end = time()
                timer.counter.update({f'{f.__module__}.{f.__name__}': end - start})
                yield i
                start = time()
        except Exception as e:
            end = time()
            timer.counter.update({f'{f.__module__}.{f.__name__}': end - start})
            raise e
        end = time()
        timer.counter.update({f'{f.__module__}.{f.__name__}': end - start})
    return inner


class AA:
    @timer
    def __init__(self):
        a = np.array(range(10_000_000))

    @generator_timer
    def __iter__(self):
        a = np.array(range(10_000_000))
        yield 'a'
        yield 'b'

    @timer
    def normal_fun(self):
        a = np.array(range(10_000_000))

    @generator_timer
    def fun_with_yield(self):
        a = np.array(range(10_000_000))
        yield a

Counter({'main.init': 1.0399727821350098, 'main.iter': 1.0183088779449463, 'main.fun_with_yield': 1.0168907642364502, 'main.normal_fun': 1.0156745910644531})