Applying decorator to all functions in a python package

Question

My question is very similar to this one , but in my case the accepted answer does not decorate all the functions in the package when they are used within the the package and I'm not sure why.

For example, I have a project set up like this:

project/
    package/ 
        __init__.py
        module_a.py
        module_b.py
    main.py

__init__.py

from .module_a import *
from .module_b import *

import types
# This is the decorator that will be used
from functools import lru_cache

for name, obj in list(globals().items()):
    if isinstance(obj, types.FunctionType):
        globals()[name] = lru_cache(maxsize=100)(obj)

module_a.py

from .module_b import b

def a(arg):
    return b

module_b.py

def b(arg):
    return arg

main.py

import package

print(package.a.cache_info())
print(package.a(None).cache_info())
print(package.b.cache_info())

When the package is imported, __init__.py decorates the functions in globals just fine when stepping through the code. However, if I execute main.py I get the following error:

C:\Users\pbreach\Anaconda3\python.exe C:/Users/pbreach/PycharmProjects/project/main.py
Traceback (most recent call last):
CacheInfo(hits=0, misses=0, maxsize=100, currsize=0)
  File "C:/Users/pbreach/PycharmProjects/project/main.py", line 4, in <module>
    print(package.a(None).cache_info())
AttributeError: 'function' object has no attribute 'cache_info'

Which would mean that b is not decorated when imported from module_b in module_a .

Why does this only happen in the second line? What might be a way to accomplish this?

I am fine with doing the decorating either during import in __init__.py or in main.py , but would rather not have to apply the decorator inside each module in package as in my case there are quite a few of them.

EDIT:

I think the issue is because globals in __init__.py is a different namespace than when b is imported into module_a meaning there are two different instances of the same function. Is there a way around this?

Answer 1

You are importing b from module_b in module_a before you get a chance to decorate it with functools.lru_cache .

The only viable alternative I see, is first explicitly decorating functions that get imported and used in other submodules and then applying the decorator to all other functions.

Using your example, first decorate b from module_b and then decorate the rest:

from package import module_b

import types
# This is the decorator that will be used
from functools import lru_cache
module_b.b = lru_cache(maxsize=100)(module_b.b)

from .module_a import *
from .module_b import *

for name, obj in list(globals().items()):
    if isinstance(obj, types.FunctionType):
        globals()[name] = lru_cache(maxsize=100)(obj)

Another option, as I stated in a comment, would be to use an if clause inside the modules that contain functions that get included in other modules were the wrapping would occur there.

So in module_b.py you could perform something like this:

if __name__ != '__main__':
    from functools import lru_cache
    b = lru_cache(b)

this only catches the case where the module isn't run as __main__ . Now when another module includes this module and it's body gets executed, the wrapping will be performed here instead of __init__.py .