In python is there a way to check if a function is a “generator function” before calling it?

后端 未结 3 1226
日久生厌
日久生厌 2020-12-13 03:44

Lets say I have two functions:

def foo():
  return \'foo\'

def bar():
  yield \'bar\'

The first one is a normal function, and the second i

相关标签:
3条回答
  • 2020-12-13 03:53
    >>> 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
    0 讨论(0)
  • 2020-12-13 04:06

    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
    
    0 讨论(0)
  • 2020-12-13 04:09
    >>> 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;-).

    0 讨论(0)
提交回复
热议问题