为所有函数、类方法及嵌套函数添加装饰器并实现模块自动导入

需求说明 #

• 为项目中所有模块的顶层函数、类方法（含同一模块内的多个类）、嵌套函数统一添加装饰器
• 装饰器需打印格式为 decorator = [函数完整名称] 的日志，确认装饰生效
• 自动扫描项目内所有文件夹的Python模块，无需手动导入

项目文件结构 #

root
├── dir1
│   └── class1.py
├── dir2
│   └── class2.py
└── main.py

原有文件代码（无需修改） #

class1.py #

class class1():
    def func1(self):
        print('class1.func1')
        def inner_func():
            print("class1.func1.inner_func")
        inner_func()
    def func2(self):
        print('class1.func2')

class2.py #

def func_from_class2():
    print("func_from_class2")
    def inner_func():
        print("func_from_class2.inner_func")
    inner_func()

class class2():
    def func1(self):
        print('class2.func1')
        def inner_func():
            print("class2.func1.inner_func")
        inner_func()
    def func2(self):
        print('class2.func2')

修改后的main.py代码 #

import inspect
import importlib
import os
import sys

# 确保项目根目录在Python路径中
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

def my_decorator(func, parent_name=""):
    """自定义装饰器，打印被装饰函数的名称"""
    def wrapper(*args, **kwargs):
        # 生成完整的函数名称（含父级函数/类的前缀）
        if parent_name:
            full_name = f"{parent_name}.{func.__name__}"
        else:
            # 处理类方法
            if inspect.ismethod(func):
                full_name = f"{func.__self__.__class__.__name__}.{func.__name__}"
            else:
                full_name = func.__name__
        print(f"decorator = {full_name}")
        return func(*args, **kwargs)
    return wrapper

def decorate_nested_functions(func, parent_name=""):
    """递归装饰函数内部的所有嵌套函数"""
    # 先装饰当前函数
    decorated_func = my_decorator(func, parent_name)
    
    # 遍历函数的常量，查找嵌套函数
    for const in func.__code__.co_consts:
        if inspect.isfunction(const):
            # 生成嵌套函数的父级名称
            nested_parent = f"{parent_name}.{func.__name__}" if parent_name else func.__name__
            # 递归装饰嵌套函数
            decorated_nested = decorate_nested_functions(const, nested_parent)
            # 替换原函数中的嵌套函数引用
            setattr(func, const.__name__, decorated_nested)
    
    return decorated_func

def decorate_class_methods(cls, class_name=""):
    """装饰类中的所有方法及方法内的嵌套函数"""
    # 遍历类的所有方法
    for name, method in inspect.getmembers(cls, predicate=lambda x: inspect.isfunction(x) or inspect.ismethod(x)):
        if not name.startswith("__"):  # 跳过内置魔法方法
            # 生成方法的父级名称（类名.方法名）
            method_parent = f"{class_name}.{name}" if class_name else name
            # 装饰方法及方法内的嵌套函数
            decorated_method = decorate_nested_functions(method, method_parent)
            setattr(cls, name, decorated_method)
    
    # 处理类内部的嵌套类
    for name, nested_cls in inspect.getmembers(cls, inspect.isclass):
        if not name.startswith("__"):
            nested_class_name = f"{class_name}.{name}" if class_name else name
            decorate_class_methods(nested_cls, nested_class_name)

def decorate_module(module):
    """装饰模块内的所有顶层函数、类及嵌套成员"""
    # 装饰顶层函数
    for name, func in inspect.getmembers(module, inspect.isfunction):
        if not name.startswith("__"):
            decorated_func = decorate_nested_functions(func)
            setattr(module, name, decorated_func)
    
    # 装饰类及类内成员
    for name, cls in inspect.getmembers(module, inspect.isclass):
        if not name.startswith("__"):
            decorate_class_methods(cls, name)

def discover_and_decorate_modules(root_dir):
    """自动发现并装饰项目内所有Python模块"""
    for root, dirs, files in os.walk(root_dir):
        # 跳过缓存目录
        if "__pycache__" in root:
            continue
        for file in files:
            if file.endswith(".py") and file != "main.py":
                # 生成模块导入路径
                module_path = os.path.relpath(os.path.join(root, file), root_dir)
                module_name = module_path.replace(os.sep, ".").replace(".py", "")
                # 导入模块
                try:
                    module = importlib.import_module(module_name)
                    # 装饰模块
                    decorate_module(module)
                    print(f"已装饰模块: {module_name}")
                except Exception as e:
                    print(f"装饰模块 {module_name} 失败: {str(e)}")

# 执行自动发现和装饰
discover_and_decorate_modules(os.path.dirname(os.path.abspath(__file__)))

# 测试代码
if __name__ == "__main__":
    # 导入模块（此时模块已被装饰）
    from dir1.class1 import class1
    from dir2.class2 import class2, func_from_class2

    print('-----------------')
    class1_obj = class1()
    class1_obj.func1()
    print('-----------------')
    class1_obj.func2()
    print('-----------------')

    class2_obj = class2()
    class2_obj.func1()
    print('-----------------')
    class2_obj.func2()
    print('-----------------')

    func_from_class2()

关键改进说明 #

1. 嵌套函数装饰：通过遍历函数的__code__.co_consts获取嵌套函数，递归调用装饰器实现深层嵌套函数的装饰
2. 自动模块发现：使用os.walk遍历项目目录，通过importlib动态导入所有非main的Python模块
3. 完整名称生成：为嵌套函数、类方法生成包含父级上下文的完整名称，便于日志追踪
4. 递归装饰逻辑：从模块→类→方法→嵌套函数逐层递归，确保所有可调用对象都被装饰

内容的提问来源于stack exchange，提问作者Fynn