如何判断一个可调用对象是否为生成器？

Question

背景

为了测试，我想追踪一个对象所有方法的调用。

现在，我写了一个循环，代码如下：

def instrument_obj(obj):
    for k in dir(obj):
        try:
            method = getattr(obj, k, None)
            if not callable(method):
                continue
            # skip reserved words:
            if k.startswith("__"):
                continue

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

这个循环对普通方法是有效的。

问题

我在代码中遇到的第一个问题是关于生成器的：

目前，my_wrapper(...)的样子是这样的：

def my_wrapper(fn):
    def _w(*args, **kwargs):
        print(f"-- called {fn.__qualname__}")
        return fn(*args, **kwargs)
    return _w

但是这样会导致生成器出问题：因为我的包装函数没有使用yield，所以调用者就无法对它进行迭代了。

对于生成器，它应该更像这样：

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

问题

我怎么能判断从我的对象中得到的可调用对象是否是生成器呢？

Answer 1

你可以使用 inspect.isgeneratorfunction 来判断一个函数是否返回了一个 generator。

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