Python function to return a list or work as a generator?

Question

I was experimenting with creating a function that returns an object or works as a generator.

This is a bad idea because, as a best practice, you want functions to reliably return the same types of values, but in the interest of science...

I'm using Python 2, and so range returns a list, and xrange is an iterable (that interestingly also provides a __len__).

def xr(start, stop=None, step=1, gen=True):
    if stop is None:
        start, stop = 0, start
    if gen == True:
        for i in xrange(start, stop, step):
            yield i
    else:
        return range(start, stop, step)

I get this error:

  File "<stdin>", line 8
SyntaxError: 'return' with argument inside generator

Questions:

Why (beyond the obvious "you can't have both yield and return in a function," assuming that's right) does it do this? Looking at my code, it is not readily apparent why it would be bad to do this.

How would I attempt to get around this? I know I can return xrange instead of yielding each item from xrange, and so I could return a generator created in another function, but is there a better way?

Answer 1

How about using generator expression?

>>> def xr(start, stop=None, step=1, gen=True):
...     if stop is None:
...         start, stop = 0, start
...     if gen == True:
...         return (i for i in xrange(start, stop, step)) # <----
...     else:
...         return range(start, stop, step)
...
>>> xr(2, gen=False)
[0, 1]
>>> xr(2, gen=True)
<generator object <genexpr> at 0x0000000002C1C828>
>>> list(xr(2, gen=True))
[0, 1]

BTW, I would rather define a generator function only. Then use list(xr(..)) if I need a list.

UPDATE Alternatively you can use iter(xrange(start, stop, step)) instead of the generator expression as @DSM commented. See Built-in functions -- iter.

Answer 2

falsetru has given you a way to have a function that returns a generator or a list. I'm answering your other question about why you can't have a return and a yield in the same function.

When you call a generator function, it doesn't actually do anything immediately (see this question). It doesn't execute the function body, but waits until you start iterating over it (or call next on it). Therefore, Python has to know if the function is a generator function or not at the beginning, when you call it, to know whether to run the function body or not.

It doesn't make sense to then have it return some value, because if it's a generator function what it returns is the generator (i.e., the thing you iterate over). If you could have a return inside a generator function, it wouldn't be reached when you called the function, so the function would have to "spontaneously" return a value at some later point (when the return statement was reached as the generator was consumed), which would either be the same as yielding a value at that time, or would be something bizarre and confusing.

I agree with falsetru that it's probably not a great idea to have a function that sometimes returns a generator and sometimes a list. Just call list on the generator if you want a list.