What is ast.arg(arg, annotation, type_comment) in Python?

Abstract Syntax Tree

ast.arg is a class defined in the ast module that expresses a single argument, in a Python list, in the form of an Abstract Syntax Tree. When the parse() method of ast is called on the Python source code that contains an argument, the ast.arg class is invoked. This class expresses the argument statement to a node in an ast tree data structure. The ast.arg class represents the arg node type in the ast tree.

Parameters

  • arg: a raw string of argument name.
  • annotation: annotation of the argument. Can be Str or Name nodes.
  • type_comment: an optional string with the type annotation as a comment.

Code

import ast
from pprint import pprint
class ArgVisitor(ast.NodeVisitor):
    def visit_arg(self, node):
        print('Node type: arg\nFields:', node._fields)
        ast.NodeVisitor.generic_visit(self, node)
    def visit_Name(self, node):
        print('Node type: Name\nFields:', node._fields)
        ast.NodeVisitor.generic_visit(self, node)
    
    def visit_FunctionDef(self, node):
        print('Node type: FunctionDef\nFields:', node._fields)
        ast.NodeVisitor.generic_visit(self, node)
    def visit_BinOp(self, node):
        print('Node type: BinOp\nFields:', node._fields)
        ast.NodeVisitor.generic_visit(self, node)
visitor = ArgVisitor()
tree = ast.parse("""
def add(a, b):
  return a + b
""")
pprint(ast.dump(tree))
visitor.visit(tree)

Explanation

  • We define an ArgVisitor class that extends from the parent class ast.NodeVisitor.
  • We override the predefined visit_arg, visit_BinOp, visit_Name, and visit_FunctionDef methods in the parent class, which receive the arg, BinOp, Name, and FunctionDef nodes, respectively.
  • In the method, we print the type and the fields inside the node and call the generic_visit() method, which invokes the propagation of visit on the children nodes of the input node.
  • We initialize a visitor object of the class ArgVisitor.
  • We write Python code that contains a function definition with its list of arguments.
  • We pass the code to the ast.parse() method, which returns the result of the expression after evaluation and store the result in a tree.
  • The ast.dump() method returns a formatted string of the tree structure in a tree.
  • The visit method available to the visitor object visits all the nodes in the tree structure.

Free Resources

License: Creative Commons-Attribution-ShareAlike 4.0 (CC-BY-SA 4.0)