I am currently trying to replace a module in a large code base in a certain condition and to figure out when any function of this module is called, I wrap each function/method in the module with a warning that prints a short message with the stack trace. The implementation of this method is as follows (printing the stack trace was omitted):
def wrap_module_with_warnings(module):
for fn_name, fn in inspect.getmembers(module):
if (not (inspect.isfunction(fn) or inspect.ismethod(fn))
or fn_name.startswith('_')
or inspect.getmodule(fn) is not module):
continue
@functools.wraps(fn)
def wrapped_fn(*args, **kwargs):
warnings.warn(f"The function {fn_name} was called.")
return fn(*args, **kwargs)
setattr(module, fn_name, wrapped_fn)
Unfortunately, this function doesn't work as I'd expect. My expectation is that it replaces all functions in module with its wrapped variant that just prints a warning and then calls the original function. However, in reality, it throws TypeError: 'module' object is not callable:
/tmp/question/main.py:22: UserWarning: The function random was called.
warnings.warn(f"The function {fn_name} was called.")
---------------------------------------------------------------------------
TypeError Traceback (most recent call last)
File /tmp/question/main.py:31
29 print(f"Output before wrapping: {mod.f()}")
30 wrap_module_with_warnings(mod)
---> 31 print(f"Output after wrapping: {mod.f()}")
File /tmp/question/main.py:23, in wrap_module_with_warnings.<locals>.wrapped_fn(*args, **kwargs)
20 @functools.wraps(fn)
21 def wrapped_fn(*args, **kwargs):
22 warnings.warn(f"The function {fn_name} was called.")
---> 23 return fn(*args, **kwargs)
TypeError: 'module' object is not callable
To demonstrate my problem, I'm using the following main code:
import mod
if __name__ == '__main__':
print(f"Output before wrapping: {mod.f()}")
wrap_module_with_warnings(mod)
print(f"Output after wrapping: {mod.f()}")
and the example module mod.py:
import random
def f() -> int:
return random.randint(0, 1000)
The error suggests that mod.f is a module. Examining the type of mod.f after calling my wrapper method reveals that it's actually a function:
In [1]: type(mod.f)
Out[1]: function
To examine the original function, I changed the last line of wrap_module_with_warnings to:
setattr(module, fn_name, (fn, wrapped_fn))
and examine mod.f after calling the wrapper doesn't help much either:
In [1]: mod.f[0]
Out[1]: <function mod.f() -> int>
In [2]: mod.f[1]
Out[2]: <function mod.f() -> int>
In [3]: type(mod.f[0])
Out[3]: function
In [4]: type(mod.f[1])
Out[4]: function
In [5]: mod.f[0]()
Out[5]: 785
In [6]: mod.f[1]()
/tmp/AscendSpeed/q/main.py:22: UserWarning: The function random was called.
warnings.warn(f"The function {fn_name} was called.")
---------------------------------------------------------------------------
TypeError Traceback (most recent call last)
Cell In[6], line 1
----> 1 mod.f[1]()
File /tmp/AscendSpeed/q/main.py:23, in wrap_module_with_warnings.<locals>.wrapped_fn(*args, **kwargs)
20 @functools.wraps(fn)
21 def wrapped_fn(*args, **kwargs):
22 warnings.warn(f"The function {fn_name} was called.")
---> 23 return fn(*args, **kwargs)
TypeError: 'module' object is not callable
I'm pretty sure that this code should work, fn should be the function mod.f, but why do I get the TypeError then?
As @jasonharper hinted at in the comments, wrapped_fn doesn't encapsulate fn_name and fn and instead uses those of the last iteration of the loop. The warning
/tmp/AscendSpeed/q/main.py:22: UserWarning: The function random was called.
warnings.warn(f"The function {fn_name} was called.")
exemplifies that. This problem is easily solved by constructing the wrapper function inside another function:
def wrap_module_with_warnings(module):
def wrap_fn(f, msg):
@functools.wraps(f)
def wrapped_fn(*args, **kwargs):
warnings.warn(msg)
return f(*args, **kwargs)
return wrapped_fn
for fn_name, fn in inspect.getmembers(module):
if (not (inspect.isfunction(fn) or inspect.ismethod(fn))
or fn_name.startswith('_')
or inspect.getmodule(fn) is not module):
continue
wrapped_fn = wrap_fn(fn, f"The function {fn_name} was called.")
setattr(module, fn_name, wrapped_fn)