Julien A
Julien A

Reputation: 111

How to make Python recognize a function as a generator function?

First, consider the following code (I'm going to discuss several versions of subgen() next):

>>> def maingen(i):
...    print("maingen started")
...    yield from subgen(i)
...    print("maingen finished")
...
>>> for i in maingen(5):
...     print(i)
...     

I want to write several subgen generator functions.

A normal one is:

>>> def subgen_1(i):
...     yield i + 1
...     yield i + 2
...     yield i + 3

No prob, output is as expected:

maingen started
6
7
8
maingen finished

Now, I want an other version of subgen which yields nothing...

If I try that:

>>> def subgen_2(i):
...     i * 3

I've an exception:

maingen started
Traceback (most recent call last):
  ...
TypeError: 'NoneType' object is not iterable

It's expected: subgen_2 isn't a generator function.

OK, next. I can read somewhere (like the first answer here) that I've to raise a StopIteration. (Note that, as a newbie, I can't comment this answer.)

>>> def subgen_3(i):
...     i * 3
...     raise StopIteration()
...     

As identified in PEP 479, "Finally, the proposal also clears up the confusion about how to terminate a generator: the proper way is return, not raise StopIteration.", I only get:

maingen started

(no maingen finished...)

And the only way I've found to get things OK is:

>>> def subgen_4(i):
...     i * 3
...     return
...     yield
... 

With this, I get:

maingen started
maingen finished

Hourrah! But this solution isn't beautiful...

Does anyone have a better or a more pythonic idea?

Is it possible to write a decorator to secretly add the ugly yield statement?

Upvotes: 2

Views: 303

Answers (2)

Julien A
Julien A

Reputation: 111

After reading all comments, I think the more pythonic way is to rewrite the maingen:

def maingen(i):
    print("maingen started")
    gen = subgen(i)
    if gen:
        yield from gen
    print("maingen finished")

With that, a subgen which is really a generator function:

def subgen_1(i):
    yield i + 1
    yield i + 2
    yield i + 3

or a subgen which is a simple function:

def subgen_2(i):
    i * 3

are both "accepted" when "injected" in the maingen generator function, and don't need some ugly yield statements.

Upvotes: 0

Tadhg McDonald-Jensen
Tadhg McDonald-Jensen

Reputation: 21453

If you want a sub generator to yield nothing, then return an empty iterator:

def subgen_2(i):
    i * 3
    return iter([]) #empty iterator

The reason you were getting TypeError: 'NoneType' object is not iterable is because without a yield statement in your sub-generator it was not a generator, but a regular function that implicitly returned None.

By returning a valid iterator that doesn't produce any values you will be able to yield from subgen_2() without error and without generating any additional values.

Another way to hide this (I don't really see why you would want to) is to make a decorator that literally just calls your function then does return iter(()) or yield from [].

def gen_nothing(f):
    @functools.wraps(f)
    def wrapper(*args,**kw):
        f(*args,**kw)
        yield from []
    return wrapper

But the only difference this produces is that the stack will require one additional frame for the wrapper, which means your Traceback messages will have a bit more noise:

Traceback (most recent call last):
  File "/Users/Tadhg/Documents/codes/test.py", line 15, in <module>
    for i in subgen():
  File "/Users/Tadhg/Documents/codes/test.py", line 6, in wrapper
    f(*args,**kw)
  File "/Users/Tadhg/Documents/codes/test.py", line 12, in subgen
    3/0
ZeroDivisionError: division by zero

Upvotes: 2

Related Questions