Friday, November 14, 2008

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

3 comments:

Jean said...

FYI,

isgen() works as advertised for Python 2.3 through 2.6.

Python 3.0rc2 throws an error due to missing attribute 'func_code'.

/Jean Brouwers

Matt Good said...

"The following flag bits are defined for co_flags: ... bit 0x20 is set if the function is a generator."

def is_generator(func):
return bool(func.func_code.co_flags & 0x20)

Ondřej Čertík said...

This is what we use in sympy (sorry for the formatting --- stupid blogger):

def isgeneratorfunction(object):
"""
Return true if the object is a user-defined generator function.

Generator function objects provides same attributes as functions.

See isfunction.__doc__ for attributes listing.

Adapted from Python 2.6.
"""
CO_GENERATOR = 0x20
if (inspect.isfunction(object) or inspect.ismethod(object)) and \
object.func_code.co_flags & CO_GENERATOR:
return True
return False