Search code examples
pythonpython-3.xgeneratordecoratorpython-decorators

Unexpected decorator behavior, "if" statement executes after if 1==2:


А few hours ago I submitted a related question and got an answer why I need to add yield to my decorator in order to function properly. I have lately recalled that I had omitted it for a reason - a strange behavior I can not explain.

I apologize in advance if I am blind, but I spent hours staring at and playing with this code and this is what I get:

def decor(func):
    def wrapper(*args, **kwargs):
        if 1==2:
            print ("Generator")
            for item in func(*args, **kwargs):
                print(item)
                #yield(item)
        else:
            print ("Not generator")
            res = func(*args, **kwargs)
            print(res)
            return res
    return wrapper

@decor            
def f():
    return "a"


f()

"""
Output:
    Not generator
    a

"""

And if I remove the comment before yield there is no output at all.

Why is that? And how is it possible that anything which I change within a if 1==2: statement takes any effect on the script?


Solution

  • If a function contains yield anywhere in the body, it is a generator function. It does not matter if the yield is executed or not. The fact that 1 == 2 is false has nothing to do with it.

    Consider the following function:

    def addone(numbers):
        for number in numbers:
            yield number + 1
    

    What happens when you call addone([])? The yield is never executed, and yet addone still returns a generator. Why should this be any different:

    def addone(numbers):
        if numbers:
            for number in numbers:
                yield number + 1
    

    So it becomes clear that whether or not the yield is actually executed is not relevant. The only relevant fact is whether the yield exists in the body of the function.

    How to fix the function

    The fix is relatively simple, all you have to do is pull the part with the yield into a separate function:

    import types
    
    def wrapper(*args, **kwargs):
        result = func(*args, **kwargs)
        if isinstance(result, types.GeneratorType):
            print("Is a generator")
            return wrap_generator(result)
        print("Not a generator")
        return result
    
    def wrap_generator(gen):
        for item in gen:
            print(item)
            yield item
    

    How to avoid in the future

    In general, the problem here is that a function is either a generator (and uses yield) or is a normal function (and uses return). It is a bit confusing when you use both yield and return in the same function!

    For Python, it turns out that if you use both yield and return in the same function, the function is a generator function. This may be somewhat confusing, so as a matter of style, I would generally avoid using both return and yield in the same function.