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

Question

Similar unanswered question: Python AST - How to see value of assign node?

I am trying to obtain the assignments statements in a file. Example:

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=[])

and form them into a dictionary like so:

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

I only want the Assigns with value and target, no call objects. The AST module is very difficult to understand and I would appreciate any help.

Answer 1

You can use ast.walk. This will recursively yield all descendant nodes. You can then check if the node is an instance of ast.Assign, if yes, it has targets and value attributes which contains the details you are looking for.

>>> 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'}