pythonpython-3.11

how can I spot if a callable is a generator?


Context

for testing purposes, I would like to trace all calls to the methods of an object.

right now: I managed to write the following loop:

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

which works for regular methods.

Problem The first issue I encountered with my code is with generators:

Right now my_wrapper(...) looks like that:

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

but this breaks generators: since my wrapper doesn't yield, a caller can't iterate on it anymore.

For generators, it should do something closer to this:

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

Question

How can I spot if the callable I got from my object is a generator ?


Solution

  • You can use inspect.isgeneratorfunction to determine if the function returns a 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