I'm having trouble applying a decorator to an imported function. Suppose I have the following foo.py module:

# contents of foo.py
def bar():
    return "hello"

I now want to import bar from it and apply the lru_cache decorator, so in my main file I do this:

from functools import lru_cache
import foo

def main():
    # caching foo.bar works fine
    for _ in range(2):
        if hasattr(foo.bar, "cache_info"):
            print("bar is already cached")
        else:
            print("caching bar")
            setattr(foo, "bar", lru_cache(foo.bar))

if __name__ == "__main__":
    main()

which produces, as expected, the following output:

caching bar
bar is already cached

However, I don't want to import the whole foo module, so now I only import bar and rely on sys.modules to find out where bar comes from. Then main.py becomes:

from functools import lru_cache
import sys
from foo import bar

def main():
    for _ in range(2):
        if hasattr(bar, "cache_info"):
            print("bar is already cached")
        else:
            print("caching bar")
            setattr(sys.modules[__name__], "bar", lru_cache(bar))


if __name__ == "__main__":
    main()

Again, this produces the wanted output:

caching bar
bar is already cached

But now I want my decorating process to be reusable in several modules, so I write it as a function in file deco.py:

from functools import lru_cache
import sys

def mydecorator(somefunction):
    for _ in range(2):
        if hasattr(somefunction, "cache_info"):
            print(f"{somefunction.__name__} is already cached")
        else:
            print(f"caching {somefunction.__name__}")
            setattr(sys.modules[somefunction.__module__], somefunction.__name__, lru_cache(somefunction))

And so main.py becomes:

import sys
from foo import bar
from deco import mydecorator

def main():
    mydecorator(bar)

if __name__ == "__main__":
    main()

This time the decoration fails; the output becomes:

caching bar
caching bar  # <- wrong; should be cached already

How do I correct my code?

The reason

When you imported bar in main.py you have implicitly assigned bar to main.py module, because

from foo import bar

actually means this:

import foo
bar = foo.bar

So, variable bar in main.py refers to original bar function.

When in deco.py you do this:

setattr(sys.modules[somefunction.__module__], "bar", lru_cache(somefunction))

you are making a NEW function object and you assign it to foo module:

1. lru_cache(somefunction) creates a new object in a memory
2. setattr(...) assign this new object to a name bar in foo module
3. However, somefunction is still the same old object (bar from main.py)

So, after you call mydecorator(bar), the situation looks like this:

1. foo.bar is now a new, wrapped func created by lru_cache(somefunction)
2. main.bar is still the same original bar() function

The solution

mydecorator() is not able to replace bar() function in-place, with a different object but keeping it under the same memory address.

What you can do, is to make mydecorator() return a new function so that you can explicitly replace bar in main.py yourself:

# deco.py

from functools import lru_cache
import sys

def mydecorator(somefunction):
    wrapped_function = None
    for _ in range(2):
        if hasattr(somefunction, "cache_info"):
            print(f"{somefunction.__name__} is already cached")
        else:
            print(f"caching {somefunction.__name__}")
            wrapped_func = lru_cache(somefunction)
            setattr(sys.modules[somefunction.__module__], somefunction.__name__, wrapped_func)
    
    return wrapped_func

# main.py

import sys
from foo import bar
from deco import mydecorator

def main():
    bar = mydecorator(bar)

if __name__ == "__main__":
    main()

This is not exactly what you wanted (you will still see "caching bar" twice), but you will get your bar decorated