Search code examples
pythonwrapperpython-decorators

How can I decorate all functions imported from a file?


I have created many functions that are divided into different files, now I would like to apply the same decorator for all of them without modifying the files and without applying the decorators one by one.

I have tried to use this explanation written by delnan, but I got no success for imported functions.

About the decorator, it must update a list every time a function within a class is executexecuted with the function arguments and values, just like this other question I asked.

Any suggestions to help me with this issue? Thanks


Solution

  • A little bit of introspection (dir()) and dynamic look-up with getattr() and setattr().

    First we iterate over all names found in module and check for objects that look like functions. After that we simply reassign old function with decorated one.

    main.py:

    import types
    import functools
    
    def decorate_all_in_module(module, decorator):
        for name in dir(module):
            obj = getattr(module, name)
            if isinstance(obj, types.FunctionType):
                setattr(module, name, decorator(obj))
    
    def my_decorator(f):
        @functools.wraps(f)
        def wrapper(*args, **kwargs):
            print(f)
            return f(*args, **kwargs)
        return wrapper
    
    import mymod1
    decorate_all_in_module(mymod1, decorator)
    

    mymod1.py:

    def f(x):
        print(x)
    
    
    def g(x, y):
        print(x + y)
    

    Output:

    <function f at 0x101e309d8>
    2
    <function g at 0x101e30a60>
    7
    

    Process does not goes that smooth if you use star imports (from mymod import *). Reason is simple - because all names are in one huge bag and there no differentiation on where they come from, you need a lot of additional tricks to find what exactly you want to patch. But, well, that's why we use namespaces - because they are one honking great idea.