如何使用Python AST获取DAG上下文管理器中tags关键字的列表值？

我正在编写一个外部工具用于分析Python源码，目标是提取出DAG上下文管理器中关键字参数tags对应的列表值。

示例源码：

from airflow.models import DAG

with DAG(
    dag_id='my_dag',
    tags=['dbt', 'marketing', 'schema', 'data_vis']
) as dag:
    pass

期望输出：

['dbt', 'marketing', 'schema', 'data_vis']

我打算用ast模块来解析文件，请问具体该如何实现？

目标源码文件（dag.py） #

from airflow.models import DAG

with DAG(
    dag_id='my_dag',
    tags=['dbt', 'marketing', 'schema', 'data_vis']
) as dag:
    pass

解析脚本（ast_explorer.py） #

import ast

with open('dag.py', 'r') as f:
    tree = ast.parse(f.read())

for node in tree.body:
    if isinstance(node, ast.With):
        for item in node.items:
            # 判断上下文表达式是否是DAG的调用
            if isinstance(item.context_expr, ast.Call) and item.context_expr.func.id == 'DAG':
                # 遍历所有关键字参数
                for arg in item.context_expr.keywords:
                    if arg.arg == 'tags':
                        # 提取列表中的字符串元素
                        tag_list = [s.s for s in arg.value.elts]

print(tag_list)
# 输出：['dbt', 'marketing', 'schema', 'data_vis']

关键逻辑说明 #

• 用ast.parse()将源码解析成抽象语法树（AST）；
• 遍历语法树中的顶级节点，筛选出With类型的节点（对应上下文管理器）；
• 在With节点的上下文表达式中，找到调用DAG函数的节点；
• 遍历DAG调用的关键字参数，匹配tags参数；
• 解析tags对应的列表节点，提取其中的字符串元素。

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