如何使用 Python AST 模块获取赋值节点的所有目标和值？

Question

类似的未回答问题：Python AST - How to see value of assign node?

我正在尝试获取文件中的赋值语句。示例：

def isMaskedArray(x): ...
isarray = isMaskedArray
masked_array = MaskedArray

_ShapeType = TypeVar("_ShapeType", bound=Any)

Module(
  body=[
    FunctionDef(
      name='isMaskedArray',     
      args=arguments(
        posonlyargs=[],
        args=[arg(
          arg='x',
          annotation=None,      
          type_comment=None)],  
        vararg=None,
        kwonlyargs=[],
        kw_defaults=[],
        kwarg=None,
        defaults=[]),
      body=[Expr(value=Constant(
        value=Ellipsis,
        kind=None))],
      decorator_list=[],
      returns=None,
      type_comment=None),
    Assign(
      targets=[Name(
        id='isarray',
        ctx=Store())],
      value=Name(
        id='isMaskedArray',
        ctx=Load()),
      type_comment=None),
    Assign(
      targets=[Name(
        id='masked_array',
        ctx=Store())],
      value=Name(
        id='MaskedArray',
        ctx=Load()),
      type_comment=None),
    Assign(
      targets=[Name(
        id='_ShapeType',
        ctx=Store())],
      value=Call(
        func=Name(
          id='TypeVar',
          ctx=Load()),
        args=[Constant(
          value='_ShapeType',
          kind=None)],
        keywords=[keyword(
          arg='bound',
          value=Name(
            id='Any',
            ctx=Load()))]),
      type_comment=None)],
  type_ignores=[])

然后像这样将它们组成字典：

[{isarray: isMaskedArray}, {masked_array: MaskedArray}]

我只想要赋值和目标，没有调用对象。 AST 模块很难理解，我将不胜感激。

Answer 1

您可以使用 ast.walk. This will recursively yield all descendant nodes. You can then check if the node is an instance of ast.Assign，如果是，它有 targets 和 value 属性，其中包含您要查找的详细信息。

>>> import ast
>>>
>>> tree = ast.parse(source)
>>>
>>> result = {
...     node.targets[0].id: node.value.id
...     for node in ast.walk(tree)
...     if isinstance(node, ast.Assign)
... }
>>>
>>> print(result)
{'isarray': 'isMaskedArray', 'masked_array': 'MaskedArray'}

Answer 2

简单 isinstance 检查独立 ast.Call 对象作为赋值值的一个问题是赋值语句的值很可能不是单个调用，而是组合调用、运算符等。因此，您可以使用递归函数来检查 ast 对象是否在任何深度包含 ast.Call 对象：

import ast
s = """
def isMaskedArray(x):
    isarray = isMaskedArray
    masked_array = MaskedArray
    ignore_val = x + y()
    ignore_val1 = val[fun()]
    a, [b, c], [[d], e, [d]] = something

_ShapeType = TypeVar("_ShapeType", bound=Any)
"""
def has_call(tree):
   return isinstance(tree, ast.Call) or any(has_call(j) for k in getattr(tree, '_fields', []) for j in 
        (getattr(tree, k) if isinstance(getattr(tree, k), list) else [getattr(tree, k)]))

r = {ast.unparse(i.targets[0]):ast.unparse(i.value) for i in ast.walk(ast.parse(s)) 
      if isinstance(i, ast.Assign) and not has_call(i)}

输出：

{'isarray': 'isMaskedArray', 'masked_array': 'MaskedArray', '(a, [b, c], [[d], e, [d]])': 'something'}

如何使用 Python AST 模块获取赋值节点的所有目标和值？

How do I use Python AST module to obtain all targets and value for assignment nodes?

python

abstract-syntax-tree