utils.py 文件源码

def get_class_node(module_node):
    class Visitor(ast.NodeVisitor):
        def visit_Module(self, node):
            self.result = None
            self.generic_visit(node)

        def visit_ClassDef(self, node):
            assert self.result is None, "More than one class in module"
            self.result = node
    v = Visitor()
    v.visit(module_node)
    return v.result