Python 从多个列表创建集合

我想计算AST中出现的变量（例如，a+b+a应该返回{a，b}）。是否有一行代码可以为每个AST案例编写？我试过：

   def count(e): 
    if isinstance(e, Add): # a + b
      # XXX can I write the following as a one liner?
      vs = set()
      for v in s.args:
        vs.update(count(v))
      return vs

    elif isinstance(e, Variable):
      return e
    ...

---

要回答实际问题，您可能需要对

Name

节点进行树遍历

这个答案包括

ast.NodeVisitor

的一个变体，它将节点的完整路径传递给访问者函数，因为我们需要读取一些上下文，以了解我们是否处于

属性

访问或

存储

正在加载或正在加载中

当然，可能有一些我还没有想到的情况，属性访问（

e.a

，

z.w

）目前被完全忽略

import ast
import collections

code = """
a = b + c + e.a
q = a + b + c
z.w = b + c
print(a)
"""


class NodeVisitorWithPath:
    def visit(self, path):
        node = path[-1]
        method = "visit_" + node.__class__.__name__
        visitor = getattr(self, method, self.generic_visit)
        return visitor(path)

    def generic_visit(self, path):
        node = path[-1]
        for field, value in ast.iter_fields(node):
            if isinstance(value, list):
                for item in value:
                    if isinstance(item, ast.AST):
                        self.visit(path + [item])
            elif isinstance(value, ast.AST):
                self.visit(path + [value])


class NameWalker(NodeVisitorWithPath):
    def __init__(self):
        self.events = collections.Counter()

    def visit_Name(self, path):
        if len(path) >= 1 and isinstance(
            path[-2], ast.Attribute
        ):  # Attribute access, ignore
            return
        event = (
            "Assign"
            if len(path) >= 1 and isinstance(path[-2], ast.Assign)
            else "Load"
        )
        node = path[-1]
        self.events[f"{event} {node.id}"] += 1
        super().generic_visit(path)


tree = ast.parse(code, "x", "exec")
nw = NameWalker()
nw.visit([tree])
print(dict(nw.events))

输出

{'Assign a': 1, 'Load b': 3, 'Load c': 3, 'Assign q': 1, 'Load a': 2, 'Load print': 1}

---

你可以取集合的并集。此外，您还可以检查

isinstance（e，set）