### 使用 AST 解析 Python 文件以提取所有 Torch API 调用 为了实现这一目标，可以利用 Python 的 `ast` 模块来解析 `.py` 文件并查找所有与 `torch` 相关的函数或类调用。下面是一个详细的解决方案： #### 定义策略 通过遍历抽象语法树 (AST)，识别出所有属于 `torch` 命名空间下的属性访问节点 (`Attribute`) 和函数调用节点 (`Call`)。 #### 实现方法 编写一段脚本用于加载待分析的目标源码文件，并构建其对应的 AST 结构；接着定义递归函数遍历整个树形结构，在遇到符合条件的节点时记录下相应的名称和位置信息。 ```python import ast def extract_torch_apis(tree): """Extract all torch related APIs from the given AST tree.""" apis = [] class TorchAPIVisitor(ast.NodeVisitor): def visit_Attribute(self, node): if isinstance(node.value, ast.Name) and node.value.id == 'torch': full_name = f'torch.{node.attr}' apis.append(full_name) elif hasattr(node.value, "attr"): # Handle nested attributes like torch.nn.Module parent_attr = getattr(node.value, "attr", "") current_attr = node.attr combined_attrs = ".".join([parent_attr, current_attr]).strip(".") if any(attr.startswith('torch.') for attr in [combined_attrs]): apis.append(combined_attrs) self.generic_visit(node) def visit_Call(self, node): if isinstance(node.func, ast.Attribute) and \ isinstance(node.func.value, ast.Name) and \ node.func.value.id == 'torch': api_call = f'{node.func.value.id}.{node.func.attr}()' apis.append(api_call) self.generic_visit(node) visitor = TorchAPIVisitor() visitor.visit(tree) return set(apis) def parse_file(file_path): with open(file_path, encoding='utf-8') as file: content = file.read() tree = ast.parse(content) return extract_torch_apis(tree) if __name__ == '__main__': result = parse_file('./path/to/your/file.py') print("\n".join(sorted(result))) ``` 此代码片段展示了如何创建自定义的 `NodeVisitor` 来处理特定类型的 AST 节点，并从中抽取有关 `torch` 库的信息[^1]。