CODEWITHSUNDEEP
pythonpython-3.xgenerator
JPI
JPI

Reputation: 113

Mixing functions and generators

So I was writing a function where a lot of processing happens in the body of a loop, and occasionally it may be of interest to the caller to have the answer to some of the computations.

Normally I would just put the results in a list and return the list, but in this case the results are too large (a few hundred MB on each loop).

I wrote this without really thinking about it, expecting Python's dynamic typing to figure things out, but the following is always created as a generator.

def mixed(is_generator=False):
    for i in range(5):
        # process some stuff including writing to a file
        if is_generator:
            yield i
    return

From this I have two questions:

1) Does the presence of the yield keyword in a scope immediately turn the object its in into a generator?

2) Is there a sensible way to obtain the behaviour I intended?

2.1) If no, what is the reasoning behind it not being possible? (In terms of how functions and generators work in Python.)

Upvotes: 1

Views: 93

Answers (1)

Netwave
Netwave

Reputation: 42698

Lets go step by step:

1) Does the presence of the yield keyword in a scope immediately turn the object its in into a generator? Yes

2) Is there a sensible way to obtain the behaviour I intended? Yes, see example below

The thing is to wrap the computation and either return a generator or a list with the data of that generator:

def mixed(is_generator=False):
    # create a generator object
    gen = (compute_stuff(i) for i in range(5))
    # if we want just the generator
    if is_generator:
        return gen
    # if not we consume it with a list and return that list
    return list(gen)

Anyway, I would say this is a bad practice. You should have it separated, usually just have the generator function and then use some logic outside:

def computation():
    for i in range(5):
        # process some stuff including writing to a file
        yield i

gen = computation()
if lazy:
    for data in gen:
        use_data(data)
else:
    data = list(gen)
    use_all_data(data)

Upvotes: 3

Related Questions