python - ast python如何在不同文件的方法之间找到连接

Question

我想在整个应用程序中索引所有方法和它们之间的连接(最终是一个包含子目录和文件的目录)。我正在使用 ast ，循环遍历目录直到单个文件，然后将它们加载到 ast像这样的对象 ast.parse(self.file_content)
我试图创建的索引是这个
connection

如果相关，这是我的代码。

def scan(self):
    '''
        scans a file line by line while keeping context of location and classes
        indexes a file buffer content into a ast, Abstract Syntax Trees, https://en.wikipedia.org/wiki/AST.
        Then, iterate over the elements, find relevant ones to index, and index them into the db.
    '''
    parsed_result = ast.parse(self.file_content)
    for element in parsed_result.body:
        results = self.index_element(element)


def index_element(self, element, class_name=None):
    '''

        if element is relevant, meaning method -> index
        if element is Class -> recursively call it self

    :param element:
    :param class_name:
    :return: [{insert_result: <db_insert_result>, 'structured_data': <method> object}, ...]
    '''
    # find classes
        # find methods inside classes
    # find hanging functions

    # validation on element type
    if self.should_index_element(element):
        if self.is_class_definition(element):
            class_element = element
            indexed_items = []
            for inner_element in class_element.body:
                # recursive call
                results = self.index_element(inner_element, class_name=class_element.name)
                indexed_items += results

            return indexed_items
        else:
            structured_data = self.structure_method_into_an_object(element, class_name=class_name)
            result_graph = self.dal_client.methods_graph.insert_or_update(structured_data)
            return "WhatEver"

    return "WhatEver"

我的问题是，是否可以使用 ast 创建此图？ .如果是，如何？
根据我的理解，我目前不能，因为我正在加载 一次一个文件 到 ast对象，它不知道外部方法。

这是我想在它们之间链接的 2 个文件的示例:
sample_a.py

from sample_class_b import SampleClassB

sample_b = SampleClassB()

class SampleClassA(object):
    def __init__(self):
        self.a = 1

    def test_call_to_another_function(self):
        return sample_b.test()

sample_b.py

class SampleClassB(object):
    def __init__(self):
        self.b = 1

    def test(self):
        return True

Answer 1

您可以遍历ast.Ast tree 并在每次递归调用时执行以下四件事之一:

如果树是 class定义，存储 class名称及其关联的方法，然后应用 Connections.walk到每个方法，存储 class和范围内的方法名称。

如果树是 import语句，加载模块并递归运行 Connections.walk在上面。

如果正在进行属性查找并且 Connections.walk在方法内，检查属性名称是否是任何 class 的方法es 当前加载。如果是这样，添加一条边到 edges将当前范围与发现的这个新方法联系起来。

如果上述情况均未发生，则继续遍历树。

import ast, itertools
import re, importlib
class Connections:
   def __init__(self):
      self._classes, self.edges = {}, []
   def walk(self, tree, scope=None):
      t_obj = None
      if isinstance(tree, ast.ClassDef):
         self._classes[tree.name] = [i for i in tree.body if isinstance(i, ast.FunctionDef) and not re.findall('__[a-z]+__', i.name)]
         _ = [self.walk(i, [tree.name, i.name]) for i in self._classes[tree.name]]
         t_obj = [i for i in tree.body if i not in self._classes[tree.name]]
      elif isinstance(tree, (ast.Import, ast.ImportFrom)):
         for p in [tree.module] if hasattr(tree, 'module') else [i.name for i in tree.names]:
            with open(importlib.import_module(p).__file__) as f:
               t_obj = ast.parse(f.read())
      elif isinstance(tree, ast.Attribute) and scope is not None:
         if (c:=[a for a, b in self._classes.items() if any(i.name == tree.attr for i in b)]):
            self.edges.append((scope, [c[0], tree.attr]))
         t_obj = tree.value
      if isinstance(t_obj:=(tree if t_obj is None else t_obj), list):
         for i in t_obj:
            self.walk(i, scope = scope)
      else:
         for i in getattr(t_obj, '_fields', []):
            self.walk(getattr(t_obj, i), scope=scope)

with open('sample_a.py') as f:
   c = Connections()
   c.walk(ast.parse(f.read()))

print(c.edges)

输出:

[(['SampleClassA', 'test_call_to_another_function'], ['SampleClassB', 'test'])] 

重要提示:取决于您正在运行的文件的复杂性 Connections.walk上，一个 RecursionError可能会发生。为了规避这一点， here是一个包含 Connections.walk 迭代版本的 Gist .

从 edges 创建图形:

import networkx as nx
import matplotlib.pyplot as plt
g, labels, c1 = nx.DiGraph(), {}, itertools.count(1)
for m1, m2 in c.edges:
   if (p1:='.'.join(m1)) not in labels:
      labels[p1] = next(c1)
   if (p2:='.'.join(m2)) not in labels:
      labels[p2] = next(c1)
   g.add_node(labels[p1])
   g.add_node(labels[p2])
   g.add_edge(labels[p1], labels[p2])
   
nx.draw(g, pos, labels={b:a for a, b in labels.items()}, with_labels = True)
plt.show() 

输出: