查找未装饰的 Python 类与函数: AST vs inspect

2025-01-06 09:40:50

查找 Python 模块中未被装饰的类和函数

在一个大量使用装饰器的 Python 代码库中，经常需要识别出哪些类或函数忘记添加装饰器。这有助于我们发现潜在的实现缺陷。本文将介绍几种方法来解决这个问题。

分析导入模块的 AST

使用 ast （抽象语法树）模块是一种相对通用的方法，可以直接分析 Python 源代码结构，获取其中的类和函数定义，然后检查它们是否被特定装饰器修饰过。该方法优点在于精确度高，但需要较多的代码实现。

操作步骤：

使用 ast.parse 将目标模块的代码解析为抽象语法树。
遍历树结构，找到所有 ClassDef 和 FunctionDef 节点。
对于每个类和函数节点，检查它的 decorator_list 属性。该属性是一个装饰器节点列表。
如果 decorator_list 为空，那么此类或函数没有被装饰，记录该结果。如果 decorator_list 非空，需继续判断其元素是否为我们预期的管理器类的注册装饰器。
打印未装饰的类和函数信息。

import ast
import os

def find_undecorated(module_path, manager_names):
    """查找给定模块中未被 manager 装饰的类和函数"""
    undecorated_elements = []
    with open(module_path, "r", encoding="utf-8") as source:
      tree = ast.parse(source.read())

    for node in ast.walk(tree):
      if isinstance(node, (ast.ClassDef, ast.FunctionDef)):
          if not node.decorator_list:
            undecorated_elements.append(
                  f"{module_path}:{node.name} is undecorated."
              )
            continue

          decorated = False
          for decorator in node.decorator_list:
            if isinstance(decorator, ast.Attribute):
              # decorator可能是像`manager.register`的形式
               if hasattr(decorator,"value") and  isinstance(decorator.value, ast.Name) and decorator.value.id in manager_names:
                    decorated = True
                    break
            elif isinstance(decorator,ast.Name):
              # decorator可能类似`@register` 或者 `my_custom_decorator`的情况, 此处未包含对于自定义decorator的处理，后续读者可以扩展此功能。
              if decorator.id in manager_names:
                  decorated=True
                  break

          if not decorated:
              undecorated_elements.append(
                  f"{module_path}:{node.name} is undecorated by any manager."
              )
    return undecorated_elements


# 测试案例
if __name__ == "__main__":
    # 示例规则文件路径
    rules_file_path = "rules.py"
    #管理器名称，即管理器类变量名或函数装饰器名
    managers_names=["manager","register"]


    # 创建测试 rules.py 文件
    with open(rules_file_path, 'w') as f:
        f.write("""
class RuleManager:
  def register(self, cls):
      return cls
manager = RuleManager()

@manager.register
class A:
    pass

class B:
    pass

@register
def C():
   pass

def D():
   pass

    """)
    undecorated = find_undecorated(rules_file_path,managers_names)
    for result in undecorated:
        print(result)
    # 清理测试文件
    os.remove(rules_file_path)

此方法提供细粒度的分析，可以精确地判断特定装饰器的存在情况。需要注意的是，上面的示例假设装饰器以属性或者名字形式出现。更为复杂的动态装饰器，以及函数装饰器将需要更多的分析。

导入模块并使用 `inspect`

另一种方法是在运行时导入模块，并使用 inspect 模块来检查模块的类和函数。 inspect 模块提供了一系列工具，用于检查模块，类，函数和其他代码对象的结构和属性。此方法简单易用，不需要额外的 AST 解析，但也存在局限性，例如它必须能够成功导入模块。

操作步骤：

使用 importlib.import_module 动态导入目标模块。
使用 inspect.getmembers 获取模块中的所有成员（包括类和函数）。
使用 inspect.isclass 和 inspect.isfunction 判断成员是否是类或函数。
使用 hasattr 以及成员对象的 __dict__ 检查是否包含特定的属性来确定是否应用了特定的装饰器。这有一定的局限性，具体实现取决于注册装饰器本身的实现方式。
记录并打印未装饰的类或函数的信息。

import inspect
import importlib.util
import os


def find_undecorated_with_inspect(module_path,manager_names):
    """
      使用inspect查找未装饰的类和函数，注意使用 inspect 时动态 import 带来的潜在风险
    """
    undecorated_elements=[]
    spec=importlib.util.spec_from_file_location("module.name",module_path)
    module =importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)


    for name,member in inspect.getmembers(module):
        if inspect.isclass(member) or inspect.isfunction(member):
            # 使用 member.__dict__ 检测是否包含指定的装饰器, 这可能不是全部的装饰器，依赖具体装饰器的实现
            is_decorated = False

            for dec_name in manager_names:
                 if hasattr(member,"__wrapped__") and  hasattr(member,"__dict__") and '__wrapped__' in member.__dict__ and member.__name__==member.__dict__['__wrapped__'].__name__ and dec_name == (member.__dict__['__wrapped__'].__name__ if hasattr(member.__dict__['__wrapped__'],'__name__') else None):
                    is_decorated=True
                    break

                 if  hasattr(member,"__dict__") and dec_name in member.__dict__:
                      is_decorated=True
                      break
                 if hasattr(member,f"__{dec_name}_dec"):#特殊情况 装饰器名称带后缀或者其他特殊名称情况，此处未包含
                      is_decorated=True
                      break


            if not is_decorated:
               undecorated_elements.append(
                f"{module_path}:{member.__name__} is undecorated by any manager"
               )
    return  undecorated_elements



# 测试案例
if __name__ == "__main__":
    # 示例规则文件路径
    rules_file_path = "rules.py"
    #管理器名称，即管理器类变量名或函数装饰器名
    managers_names=["register"]


    # 创建测试 rules.py 文件
    with open(rules_file_path, 'w') as f:
        f.write("""
class RuleManager:
  def register(self, cls):
      return cls
manager = RuleManager()

@manager.register
class A:
    pass

class B:
    pass
@register
def C():
   pass

def D():
   pass
    """)

    undecorated=find_undecorated_with_inspect(rules_file_path,managers_names)
    for result in undecorated:
         print(result)


    # 清理测试文件
    os.remove(rules_file_path)

inspect 模块提供的运行时检查功能方便快捷。但依赖 import 和 __dict__ 可能存在缺陷，对于复杂的动态装饰器识别能力稍差。此外，运行时导入执行代码也有安全风险，应该避免在不受信任的模块上执行此操作。

安全建议

使用 AST 分析方法能够获得最精确的结果，它不依赖于运行时行为，并且相对安全，适合用于静态代码分析工具。使用 inspect 方法的时候需要特别注意运行导入的代码来源是否可信，并尽可能使用 importlib 中相关的模块来实现动态 import。另外使用代码检测工具，例如 pylint, flake8 和 mypy 等可以有效地降低引入问题的风险。

选择何种方式取决于项目环境的需求和安全性考量。