ast获取指定python文件中携带指定装饰器函数的实现

在实现自动化测试过程中，需要根据指定的装饰器来标记需要执行的用例或函数，下面根据使用ast库来实现读取指定文件中的数据结构后对其内容进行解析并拿到携带对应装饰器的函数。
根据以下方法仅能解析func、class-func的数据结构，其余数据结构可能不兼容，需要根据实际情况进行完善调整。

# 被解析文件
import pytest

from AutoTestMain.common.MyDecorator import sj_auto
from libs.sj_logger import logger


class Test01:
    @sj_auto(case_name="自动化用例1")
    @pytest.mark.smoke
    def test_01(self, args, kwargs):
        raise ValueError("test1")
        print("执行了类函数用例1")

    @sj_auto(case_name="自动化用例7")
    @pytest.mark.smoke
    def test_07(self, args, kwargs):
        print("执行了类函数用例7")


def setup():
    print("模块内所有用例执行前需要执行的内容")


def teardown():
    print("模块内所有用例执行完毕后需要执行的内容")


class Test02:

    def setup_class(self):
        print("在类内所有案例执行前执行内容")

    def teardown_class(self):
        print("在类内所有案例执行完成后执行的内容")

    def setup_method(self):
        print("在执行用例前执行的代码")

    def teardown_method(self):
        print("在执行用例后执行的代码")

    @sj_auto(case_name="自动化用例2", case_level="P0")
    @pytest.mark.smoke
    def test_02(self, test1: str, except_result):
        logger.info(f"执行了用例：{test1}")
        logger.info(f"对比了预期：{except_result}")

    @sj_auto(case_name="自动化用例4", case_level="P0")
    @pytest.mark.smoke
    def test_04(self, test1: str, except_result):
        logger.info(f"执行了用例：{test1}")
        logger.info(f"对比了预期：{except_result}")

    @sj_auto(case_name="自动化用例5", case_level="P0")
    @pytest.mark.smoke
    def test_05(self, test1: str, except_result):
        logger.info(f"执行了用例：{test1}")
        logger.info(f"对比了预期：{except_result}")


@sj_auto(case_name="自动化用例3", case_level="P1")
@pytest.mark.smoke
@pytest.mark.skip(reason="不想执行")
def test_03(test1: str, test2: int):
    """
    测试使用的用例
    :param test1:
    :param test2:
    :return:
    """
    logger.info("执行了用例3")


@sj_auto()
@pytest.mark.smoke
@pytest.mark.skip(reason="不想执行")
def test_06(test1: str, test2: int):
    """
    测试使用的用例
    :param test1:
    :param test2:
    :return:
    """
    logger.info("执行了用例6")


@pytest.mark.smoke
@pytest.mark.usermanage
def test_04():
    print("执行了用例1")

# 解析函数
import ast

from libs.sj_logger import logger


def find_decorated_functions(file_path, decorator_name):
    with open(file_path, 'r', encoding='utf-8') as f:
        source_code = f.read()
    tree = ast.parse(source_code)
    decorated_functions = []
    tree_parse(tree, decorated_functions, decorator_name)
    logger.info(f"获取到指定装饰器目录结构：{decorated_functions}")
    return decorated_functions


def tree_parse(tree, decorated_functions, decorator_name, class_name=None, class_check=0):
    """当前仅处理func、class-func类型的数据，如有使用其他层级结构的情况，当前会处理异常"""
    for node in tree.body:
        if isinstance(node, ast.FunctionDef):
            if any((isinstance(decorator, ast.Name) and decorator.id == decorator_name) or (
                    isinstance(decorator, ast.Call) and decorator.func.id == decorator_name) for
                   decorator in node.decorator_list):
                node_info = {"class_name": None, "funcs_name": []}
                if class_name is not None:
                    class_check += 1
                    node_info["class_name"] = class_name
                # 如果是0，则说明不是class结构，则直接添加到装饰器目录结构
                if class_check == 0:
                    node_info["funcs_name"].append(node.name)
                    decorated_functions.append(node_info)
                else:
                    # 说明是第一次进入一个类，直接将对应数据记录到装饰器目录结构
                    if class_check == 1:
                        node_info["funcs_name"].append(node.name)
                        decorated_functions.append(node_info)
                    # 说明是多次进入一个类，需要从已有的装饰器目录结构中拿到同类的数据
                    else:
                        for node_tree in decorated_functions:
                            if node_tree["class_name"] == class_name:
                                node_tree["funcs_name"].append(node.name)
                                break
        elif isinstance(node, ast.ClassDef):
            class_name = node.name
            class_check += 1
            tree_parse(node, decorated_functions, decorator_name, class_name)
            # 说明一个class已经处理完成，需要还原class_name
            class_check = 0
            class_name = None
        else:
            # 其他的可能是import之类的，无需处理
            pass


if __name__ == '__main__':
    case_path = r"D:\MyCode\AjTmp\backend\AutoTestMain\PyCase\test_demo.py"
    functions = find_decorated_functions(case_path, decorator_name="sj_auto")
    # print(functions)

执行结果：

[
    {
        "class_name": "Test01",
        "funcs_name": [
            "test_01",
            "test_07"
        ]
    },
    {
        "class_name": "Test02",
        "funcs_name": [
            "test_02",
            "test_04",
            "test_05"
        ]
    },
    {
        "class_name": null,
        "funcs_name": [
            "test_03"
        ]
    },
    {
        "class_name": null,
        "funcs_name": [
            "test_06"
        ]
    }
]