如何用AST分析Python仓库，区分三类调用并提取内部函数调用

针对你从GitHub仓库提取内部函数调用的需求，以下是一套基于AST的高效实现流程，以及工具替代方案，同时附带其他语言的思路：

一、Python 核心实现步骤 #

1. 预收集仓库内所有本地符号 #

先遍历整个仓库的.py文件，用AST提取每个文件中定义的函数、类方法，建立本地符号库：

import os
import ast

local_modules = set()
local_symbols = {}  # 格式：{模块名: {函数名集合}}

def collect_local_symbols(root_dir):
    for root, _, files in os.walk(root_dir):
        for file in files:
            if file.endswith('.py'):
                file_path = os.path.join(root, file)
                # 生成模块名（比如src/utils.py → src.utils）
                rel_path = os.path.relpath(file_path, root_dir)
                module_name = os.path.splitext(rel_path)[0].replace(os.sep, '.')
                local_modules.add(module_name)
                
                with open(file_path, 'r', encoding='utf-8') as f:
                    tree = ast.parse(f.read())
                
                # 提取函数和类方法
                funcs = set()
                for node in ast.walk(tree):
                    if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                        funcs.add(node.name)
                    elif isinstance(node, ast.ClassDef):
                        for item in node.body:
                            if isinstance(item, (ast.FunctionDef, ast.AsyncFunctionDef)):
                                funcs.add(item.name)
                local_symbols[module_name] = funcs

# 调用：传入仓库根目录
collect_local_symbols('/path/to/repo')

2. 解析导入语句，建立符号映射 #

对每个待分析的文件，解析其导入语句，记录每个局部符号的来源类型（本地/第三方/标准库）：

import sys

def build_import_mapping(tree, current_module):
    import_map = {}  # 格式：{局部符号: (来源类型, 来源模块)}
    
    for node in ast.walk(tree):
        # 处理import ... as ...
        if isinstance(node, ast.Import):
            for alias in node.names:
                mod_name = alias.name
                local_name = alias.asname or mod_name
                if mod_name in sys.stdlib_module_names:
                    import_map[local_name] = ('stdlib', mod_name)
                elif mod_name in local_modules:
                    import_map[local_name] = ('local', mod_name)
                else:
                    import_map[local_name] = ('third-party', mod_name)
        
        # 处理from ... import ... as ...
        elif isinstance(node, ast.ImportFrom):
            mod_name = node.module or current_module  # 处理相对导入
            # 解析相对导入（比如from .utils import ... → 当前模块.parent.utils）
            if node.level > 0:
                parts = current_module.split('.')
                mod_name = '.'.join(parts[:-node.level]) + ('.' + mod_name if mod_name else '')
            
            for alias in node.names:
                local_name = alias.asname or alias.name
                if mod_name in sys.stdlib_module_names:
                    import_map[local_name] = ('stdlib', mod_name)
                elif mod_name in local_modules:
                    # 检查符号是否在本地模块的符号库中
                    if alias.name in local_symbols.get(mod_name, set()):
                        import_map[local_name] = ('local', mod_name)
                    else:
                        import_map[local_name] = ('third-party', mod_name)
                else:
                    import_map[local_name] = ('third-party', mod_name)
    
    return import_map

3. 分析函数调用，筛选内部调用 #

遍历AST的Call节点，结合导入映射判断调用是否来自本地：

def extract_local_calls(tree, import_map):
    local_calls = []
    
    for node in ast.walk(tree):
        if isinstance(node, ast.Call):
            # 处理直接调用（比如abc()）
            if isinstance(node.func, ast.Name):
                symbol = node.func.id
                if symbol in import_map and import_map[symbol][0] == 'local':
                    local_calls.append({
                        'function': symbol,
                        'module': import_map[symbol][1],
                        'line': node.lineno
                    })
            # 处理模块调用（比如file_b.abc()）
            elif isinstance(node.func, ast.Attribute):
                # 回溯属性链的顶层符号
                top_node = node.func
                while isinstance(top_node, ast.Attribute):
                    top_node = top_node.value
                if isinstance(top_node, ast.Name):
                    symbol = top_node.id
                    if symbol in import_map and import_map[symbol][0] == 'local':
                        full_func_name = f"{symbol}.{node.func.attr}"
                        local_calls.append({
                            'function': full_func_name,
                            'module': import_map[symbol][1],
                            'line': node.lineno
                        })
    
    return local_calls

4. 整合流程 #

def process_repo_file(file_path, root_dir):
    with open(file_path, 'r', encoding='utf-8') as f:
        tree = ast.parse(f.read())
    
    # 获取当前文件的模块名
    rel_path = os.path.relpath(file_path, root_dir)
    current_module = os.path.splitext(rel_path)[0].replace(os.sep, '.')
    
    import_map = build_import_mapping(tree, current_module)
    local_calls = extract_local_calls(tree, import_map)
    return local_calls

# 示例：处理file_a.py
calls = process_repo_file('/path/to/repo/file_a.py', '/path/to/repo')
for call in calls:
    print(f"内部调用：{call['function']}（来自模块{call['module']}，行{call['line']}）")

二、工具替代方案 #

如果不想手动写AST逻辑，可以用更智能的静态分析库：

• astroid：比原生AST更强大，支持符号推断，能直接获取函数的定义位置。比如用node.func.infer()推断符号来源，判断是否属于当前项目目录。
• pyright：微软的Python类型检查器，可通过其API获取符号的导入来源，适合大规模项目分析。

三、其他语言的实现思路 #

• JavaScript/TypeScript：用typescript的Compiler API或esprima解析AST，遍历项目文件收集导出的函数，分析import语句的路径，判断是否为项目内相对路径。
• Java：用JavaParser解析.java文件，收集类和方法，对比import语句的包名是否属于项目的包结构。
• Go：用go/ast和go/packages包，遍历项目文件收集函数，分析import路径是否为本地包（相对路径或项目模块名）。

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