Wrapping a function imported with import *?

Question

I have a set of tests which wrap some functions of a specific module suppose it looks like this: alpha.py

def foo(): 
    return 'foo'
def bar(): 
    return foo(), 'bar'

debug_alpha.py

from alpha import *
def _wrapper(func): 
    def call(*args, **kwargs):
        print('entering')
        result = func(*args, **kwargs)
        print('exiting') 
        return result
    return call

def _wrapFunc(module, func): 
    module.__dict__[func.__name__] = _wrapper(func)

test.py

import debug_alpha
debug_alpha._wrapFunc(debug_alpha, debug_alpha.foo)
print(debug_alpha.foo())
print(debug_alpha.bar())

But of course the output is only:

entering
exiting
foo
('foo', 'bar')

Because while the foo function may have been wrapped - bar is still referenced to the foo function in the original alpha module. Is there a trick to wrapping foo in a way that bar will also call the appropriate foo without changing the import schema?

Answer 1

I'm not entirely sure of your objective, but you could do something similar to the following:

alpha.py

# everything here is the same
def foo(): 
    return 'foo'

def bar(): 
    return foo(), 'bar'

debug.py Use the inspect module to find the original module that the function was defined in.

# make the debug module reusable
import inspect

def _wrapper(func): 
    def call(*args, **kwargs):
        print('entering')
        result = func(*args, **kwargs)
        print('exiting') 
        return result
    return call


def _wrapFunc(func): 
    # find out which module the original function was
    # declared and wrap it.
    module = inspect.getmodule(func)
    wrapped = _wrapper(func)
    setattr(module, func.__name__, wrapped)
    # return the wrapped function so that the module
    # that called this function can have access to the
    # newly created function
    return wrapped

test_alpha.py

# make test.py the testing script instead of depending on
# debug_alpha to handle hardcoded namespaces
import debug_alpha
from alpha import *

# wrap the function in its original module
# and override the wrapped function in this namespace
foo = debug_alpha._wrapFunc(foo)

print(foo())
print(bar())

Output

entering
exiting
foo
entering
exiting
('foo', 'bar')