如何通过Python AST确定函数调用所属的模块？

Python AST没有内置方法直接确定函数调用所属的模块，需要自己实现逻辑来跟踪导入的符号映射，并解析函数调用的完整路径。

核心思路 #

要找出datetime.now()这类调用所属的模块，需要完成两件核心工作：

• 遍历AST收集所有导入语句，建立本地符号到模块路径的映射表
• 遍历AST中的函数调用节点，解析调用的完整限定名，再和映射表匹配

1. 收集导入映射 #

针对两种常见导入方式分别处理：

• 普通import：比如import datetime，本地符号datetime直接对应根模块datetime；如果是import datetime as dt，则dt对应datetime
• from ... import ...：比如from datetime import datetime，本地符号datetime对应datetime.datetime；如果是from datetime import datetime as dt，则dt对应datetime.datetime

2. 解析函数调用的限定名 #

对于函数调用节点（Call），需要向上追溯它的func部分：

• 如果是Attribute节点（比如datetime.now()中的datetime.now），递归拼接属性链得到完整路径
• 如果是Name节点，直接取符号名，再通过映射表转换为模块路径

代码实现 #

以下是一个AST Visitor类的示例，可实现上述逻辑：

import ast
from typing import Dict, List

class CallTracker(ast.NodeVisitor):
    def __init__(self):
        self.import_map: Dict[str, str] = {}
        self.target_calls: List[str] = []  # 存储匹配到的目标调用路径

    def visit_Import(self, node: ast.Import):
        for alias in node.names:
            # 处理 import [模块] [as 别名] 的场景
            self.import_map[alias.asname or alias.name] = alias.name
        self.generic_visit(node)

    def visit_ImportFrom(self, node: ast.ImportFrom):
        module_base = node.module or ""
        for alias in node.names:
            # 处理 from [模块] import [对象] [as 别名] 的场景
            full_path = f"{module_base}.{alias.name}" if module_base else alias.name
            self.import_map[alias.asname or alias.name] = full_path
        self.generic_visit(node)

    def _get_full_path(self, node) -> str:
        """递归解析节点的完整限定路径"""
        if isinstance(node, ast.Name):
            return self.import_map.get(node.id, node.id)
        elif isinstance(node, ast.Attribute):
            parent_path = self._get_full_path(node.value)
            return f"{parent_path}.{node.attr}"
        else:
            return ""

    def visit_Call(self, node: ast.Call):
        full_path = self._get_full_path(node.func)
        # 匹配 datetime.now() 和 datetime.datetime.now() 两种场景
        if full_path in ("datetime.datetime.now", "datetime.now"):
            self.target_calls.append(full_path)
        self.generic_visit(node)

# 测试普通import场景
code1 = """
import datetime
datetime.datetime.now()
"""
tracker1 = CallTracker()
tracker1.visit(ast.parse(code1))
print("普通import场景匹配结果：", tracker1.target_calls)  # 输出 ['datetime.datetime.now']

# 测试from import场景
code2 = """
from datetime import datetime
datetime.now()
"""
tracker2 = CallTracker()
tracker2.visit(ast.parse(code2))
print("from import场景匹配结果：", tracker2.target_calls)  # 输出 ['datetime.now']

说明 #

• 该Visitor会先遍历所有导入语句构建符号映射表，再遍历函数调用节点解析完整路径
• 可根据需求扩展逻辑，比如支持嵌套模块导入、批量匹配多个目标函数等

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