73

Lets say I have two functions:

def foo():
  return 'foo'

def bar():
  yield 'bar'

The first one is a normal function, and the second is a generator function. Now I want to write something like this:

def run(func):
  if is_generator_function(func):
     gen = func()
     gen.next()
     #... run the generator ...
  else:
     func()

What will a straightforward implementation of is_generator_function() look like? Using the types package I can test if gen is a generator, but I wish to do so before invoking func().

Now consider the following case:

def goo():
  if False:
     yield
  else:
     return

An invocation of goo() will return a generator. I presume that the python parser knows that the goo() function has a yield statement, and I wonder if it possible to get that information easily.

Thanks!

Carlos
  • 932
  • 1
  • 7
  • 9
  • 3
    It's useful to note that if a function contains a `yield` statement, then a `return` statement inside that function is not permitted to have an argument. It has to be just `return` which terminates the generator. Good question! – Greg Hewgill Dec 09 '09 at 05:13
  • Good point, `goo()` should not be valid, however it is, at least here (Python 2.6.2). – Carlos Dec 09 '09 at 05:32
  • 7
    A note to current readers: @GregHewgill comment above is no longer right, now you can return with argument (which is passed on the value attr of the StopIteration) – wim Mar 14 '14 at 01:58

3 Answers3

88
>>> import inspect
>>> 
>>> def foo():
...   return 'foo'
... 
>>> def bar():
...   yield 'bar'
... 
>>> print inspect.isgeneratorfunction(foo)
False
>>> print inspect.isgeneratorfunction(bar)
True
  • New in Python version 2.6
Corey Goldberg
  • 59,062
  • 28
  • 129
  • 143
  • 51
    Just a 2014 comment thanking you for providing a 2011 answer on a 2009 question :) – wim Mar 14 '14 at 02:07
  • 4
    Thought this solved my problem, but it doesn't perfectly. If the function is a wrapped generator, such as `partial(generator_fn, somearg=somevalue)` then this won't be detected. Nor will a lambda used in a similar circumstance, such as `lambda x: generator_fun(x, somearg=somevalue)`. These actually work as expected; he code was experimenting with a helper function that can chain generators, but if a normal function is found, it'll wrap it in a "single item generator". – Chris Cogdon Aug 25 '16 at 18:38
  • 6
    Just a 2020 comment appreciating that in 2014 you thanked him for providing a 2011 answer on a 2009 question. – Vitalate Dec 19 '20 at 13:25
  • Note that `inspect.isgeneratorfunction()` doesn't work for AsyncGenerator. Maybe time did a trick on that answer ? ;) – Cyril N. Dec 01 '21 at 09:57
  • @CyrilN. `inspect.isasyncgenfunction()` exists for that. [docs](https://docs.python.org/3/library/inspect.html#inspect.isasyncgenfunction) – lovre Jan 31 '23 at 15:54
7
>>> def foo():
...   return 'foo'
... 
>>> def bar():
...   yield 'bar'
... 
>>> import dis
>>> dis.dis(foo)
  2           0 LOAD_CONST               1 ('foo')
              3 RETURN_VALUE        
>>> dis.dis(bar)
  2           0 LOAD_CONST               1 ('bar')
              3 YIELD_VALUE         
              4 POP_TOP             
              5 LOAD_CONST               0 (None)
              8 RETURN_VALUE        
>>> 

As you see, the key difference is that the bytecode for bar will contain at least one YIELD_VALUE opcode. I recommend using the dis module (redirecting its output to a StringIO instance and checking its getvalue, of course) because this provides you a measure of robustness over bytecode changes -- the exact numeric values of the opcodes will change, but the disassembled symbolic value will stay pretty stable;-).

Alex Martelli
  • 854,459
  • 170
  • 1,222
  • 1,395
  • Alex, how do you feel about calling "blah = func()"... then checking if type(blah) is a generator? and if it's not, then func() was called already :-). I think that would have been how I would have first investigated how to do this :-). – Tom Dec 09 '09 at 05:13
  • Was about to write the same but the Python übergod came in first. :-) – paprika Dec 09 '09 at 05:14
  • 1
    The OP is very clear in the Q's title that he wants the information **before calling** -- showing how to get it **after** calling does not answer the given question with the clearly expressed constraints. – Alex Martelli Dec 09 '09 at 05:21
  • @paprika: Ha... I have no idea if this works... but Alex said it does... so +1 :-)... and I doubt anyone else will have a better answer... not even Guido himself. – Tom Dec 09 '09 at 05:22
  • @Alex: absolutely... I must say... I didn't read the question thoroughly :-o. Now I'm looking at it and it's obvious :-). I kind of focused on the "run" function and thought that could easily be done with types. – Tom Dec 09 '09 at 05:25
  • @paprika, tx for the kudos, but that title belongs to Guido van Rossum (who's not active on SO, AFAIK). In the Python pantheon, I can at most be an Einherjar (gender issues precluding me being a Valkyrie;-). – Alex Martelli Dec 09 '09 at 05:29
1

I've implemented a decorator that hooks on the decorated function returned/yielded value. Its basic goes:

import types
def output(notifier):
    def decorator(f):
        def wrapped(*args, **kwargs):
            r = f(*args, **kwargs)
            if type(r) is types.GeneratorType:
                for item in r:
                    # do something
                    yield item
            else:
                # do something
                return r
    return decorator

It works because the decorator function is unconditionnaly called: it is the return value that is tested.


EDIT: Following the comment by Robert Lujo, I ended up with something like:

def middleman(f):
    def return_result(r):
        return r
    def yield_result(r):
        for i in r:
            yield i
    def decorator(*a, **kwa):
        if inspect.isgeneratorfunction(f):
            return yield_result(f(*a, **kwa))
        else:
            return return_result(f(*a, **kwa))
    return decorator
Damien
  • 1,624
  • 2
  • 19
  • 26
  • I had similar case and I got error: SyntaxError: 'return' with argument inside generator. When I think about it, it looks logical, the same function can't be normal function and generator function in the same time. Does this really work in your case? – Robert Lujo Sep 03 '15 at 16:04