Sometimes I want a generator-CodePudding

The following code took me by surprise. I was hoping that I could write a function that might or might not act as a generator.

def mayGen(limit=5, asGenerator=False):
    result = []
    for i in range(1, limit   1):
        n = i * i
        if asGenerator:
            yield n
        else:
            result.append(n)
    return result

print(mayGen(5, False))

for x in mayGen(5, True):
    print(x)

But no, the output of this program is

<generator object mayGen at 0x7fa57b6ea7b0>
1
4
9
16
25

Calling mayGen with asGenerator=False is simply useless. It seems that the mere existence of a yield statement, regardless of whether it is executed, radically changes the behavior of the containing function.

So what if mayGen was actually big and complicated and I wish I could ask it to either function as a generator, or write its computed output to a file? Is there a well-traveled path here? How are people handling this sort of thing?

CodePudding user response：

Just write a generator. You can always take its output and do whatever you want with it: put it in a list, write it to a file, etc. Although, if that's what you're going to be using it for a lot, you might want to write a wrapper function to do it for you, for example:

def willGen(limit=5):
    for i in range(1, limit 1):
        n = i * i
        yield n

def willReturnList(*args, **kwargs):
    return list(willGen(*args, **kwargs))

CodePudding user response：

Edit: For a number of reasons mentioned in comments, this is probably not a great idea. In any event, you can write this:

def mayGen(limit=5, asGenerator=False):
    def _generate(limit):
        for i in range(1, limit   1):
            n = i * i
            yield n
    def _return(limit):
        result = []
        for i in range(1, limit   1):
            n = i * i
            result.append(n)
        return result
    if asGenerator:
        return _generate(limit)
    return _return(limit)

Edit: to simiplify even further, you should probably just return list(_generate(limit)) in the final line.