Python: introspecting for generator vs function

At Baypiggies yesterday evening, Fernando Perez asked how to tell by introspection whether you're dealing with a "plain old Pythjon-coded function", like, say,

def f(): return 1

or rather a generator function, like, say,

def y(): yield 1

Apparently, he needs that info to perfect some decorator they're using as part of a nose-based test framework for scipy.

I don't think there's any other way except by looking at the bytecode for Yield vs Return opcodes; so, I jotted down on the spot the quick-and-dirty approach based on that idea:

import dis, sys
def is_generator(f):
   save_stdout = sys.stdout
   fake_stdout = cStringIO.StringIO()
   try:
       sys.stdout = fake_stdout
       dis.dis(f)
   finally:
       sys.stdout = save_stdout
   return ' YIELD_VALUE ' in fake_stdout.getvalue()

Of course, this does much more work than necessary (AND can be fooled by a function containing a peculiar string literal... like itself!-) . So, early this morning, I prepped a somewhat better performing and more solid version:

import opcode
YIELD_OP = opcode.opmap['YIELD_VALUE']
RETURN_OP = opcode.opmap['RETURN_VALUE']
HAVE_ARG = opcode.HAVE_ARGUMENT

def isgen(f):
    code = f.func_code.co_code
    n = len(code)
    i = 0
    while i < n:
        op = ord(code[i])
        i += 1
        if op >= HAVE_ARG:
            i += 2
        elif op == YIELD_OP:
            return True
        elif op == RETURN_OP:
            return False
    return False

Hope this can prove useful to someone else!-)


Alex