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 ?
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