Search code examples
pythongeneratordecorator

Decorated generator function


I have a decorator:

def remediation_decorator(dec_mthd):
    def new_func(*args, **kwargs):
        try:
            return dec_mthd(*args, **kwargs)
        except (KeyError, HTTPError) as err:
            print(f'error = {err}... call the remediation function')
    return new_func

Inside the generator function, another function is called to raise specific exceptions under certain conditions:

def check(number):
    if number == 1:
        raise HTTPError
    if number == 2:
        raise KeyError

This generator function is decorated like so:

@remediation_decorator
def dec_mthd_b(number):
    check(number)
    for i in range(0,3):
        yield i+1

When an exception is raised by the check function, the decorator's except is not hit.

[ins] In [16]: dec_mthd_b(1)
Out[16]: <generator object dec_mthd_b at 0x10e79cc80>

It appears to behave like this because it's a generator function - from Yield expressions:

When a generator function is called, it returns an iterator known as a generator.

(I wonder whether to take this in the literal sense 'it returns the iterator first irrespective of other logic in the function', hence why check() does not raise the exception?)

and,

By suspended, we mean that all local state is retained, including the current bindings of local variables, the instruction pointer, the internal evaluation stack, and the state of any exception handling.

Have I understood this correctly? Please can anyone explain this further?


Solution

  • Yes you got it. @remediation_decorator is a Syntactic Sugar in python for decorators. I'm going to use the verbose(?) form:

    def dec_mthd_b(number):
        check(number)
        for i in range(0, 3):
            yield i + 1
    
    dec_mthd_b = remediation_decorator(dec_mthd_b)
    

    What does this line do ? remediation_decorator is your decorator, it gives you the inner function, in your case new_func.

    What is new_func ? It is a normal function, when you call it, it runs the body of the function.

    What will return from new_func ? dec_mthd(*args, **kwargs).

    Here dec_mthd points to dec_mthd_b and it is a function again. But when you call it, since dec_mthd_b has yield` keyword inside, it gives you back the generator object.

    Now here is the point. The body of your inner function, here new_func, is executed without any problem. You got your generator object back. No error is raised...

    # this is the result of calling inner function, which gives you the generator object
    gen = dec_mthd_b(1)
    
    # Here is where you're going to face the exceptions.
    for i in gen:
        print(i)
    

    What will happen in the for loop ? Python runs the body of the dec_mthd_b. The error is raised from there...

    So in order to catch the exceptions, you have two options, either catch it inside the dec_mthd_b, or in the last for loop.