在Python对象方法插桩中，如何判断可调用对象是否为生成器？

背景 #

出于测试需求，我需要追踪某对象的所有方法调用，已经编写了遍历对象方法并进行包装的代码：

def instrument_obj(obj):
    for k in dir(obj):
        try:
            method = getattr(obj, k, None)
            if not callable(method):
                continue
            # 跳过内置特殊方法
            if k.startswith("__"):
                continue

            wrapped = my_wrapper(method)
            setattr(obj, k, wrapped)
            print(f"-- instrumented {k}")
        except:
            pass

这段代码对常规方法有效，但遇到生成器方法时会出现问题：普通的包装器直接返回生成器对象，无法保持生成器的迭代特性——因为包装器没有使用yield，导致调用者无法正常迭代。

针对生成器，需要改用带yield的包装逻辑：

def my_wrapper(gen):
    def _w(*args, **kwargs):
        print(f"-- called {gen.__qualname__}")
        for itm in gen(*args, **kwargs):
            yield itm
    return _w

核心问题 #

如何判断从对象中获取的callable是否为生成器函数，从而选择对应的包装逻辑？

解决方案 #

方法1：使用标准库inspect的isgeneratorfunction（推荐） #

Python的inspect模块提供了专门的函数来判断一个函数是否是生成器函数，这是最直观且跨版本可靠的方法：

import inspect

def my_wrapper(fn):
    # 判断是否为生成器函数
    if inspect.isgeneratorfunction(fn):
        def _wrapped_generator(*args, **kwargs):
            print(f"-- called {fn.__qualname__}")
            # 用yield from简化生成器迭代转发
            yield from fn(*args, **kwargs)
        return _wrapped_generator
    else:
        def _wrapped_function(*args, **kwargs):
            print(f"-- called {fn.__qualname__}")
            return fn(*args, **kwargs)
        return _wrapped_function

inspect.isgeneratorfunction()会检查目标函数的字节码是否包含生成器特性（即是否使用了yield或yield from），返回布尔值，完美适配我们的需求。

方法2：检查函数的底层标志位（不推荐，依赖CPython实现） #

如果不想引入inspect模块，可以直接检查函数对象的__code__.co_flags是否包含CO_GENERATOR标志位（对应十六进制0x20）：

import types

def is_generator_function(fn):
    # 先判断是否是函数对象，再检查标志位
    return isinstance(fn, types.FunctionType) and (fn.__code__.co_flags & 0x20) != 0

# 修改my_wrapper使用这个判断
def my_wrapper(fn):
    if is_generator_function(fn):
        def _wrapped_generator(*args, **kwargs):
            print(f"-- called {fn.__qualname__}")
            yield from fn(*args, **kwargs)
        return _wrapped_generator
    else:
        def _wrapped_function(*args, **kwargs):
            print(f"-- called {fn.__qualname__}")
            return fn(*args, **kwargs)
        return _wrapped_function

这种方法依赖CPython的内部实现细节，不同Python解释器（如PyPy）可能有差异，因此优先推荐使用inspect模块的方法。

内容的提问来源于stack exchange，提问作者LeGEC