Python AST

You will be going to learn how to use AST (abstract syntax trees) for your Python source code.

Before moving to it, you should have an understanding of python interpreter.

How Python interpreter works:

The first code is parsed into the list of tokens. Now, this list of tokens is used to build an abstract syntax tree (collection of nodes which are linked together). The interpreter generates the bytecode from the tree, then this bytecode is used to run your python source code.

What is AST:

Python has “AST” module library.

AST (abstract syntax tree) provides tree structure for your source code (Python or any other programming language). In Python, AST is compiled into a code object using the built-in “compile()” function.

Explanation using tree diagram:

Abstract Syntax tree for expression:

AST tree for expression: 5+7*3

Python abstract syntax tree structure: x = y + 3

Where Assign and BinOp is AST node type.

AST classes and fields:

• AST: Node abstract class
• _fields: To get all child nodes
• lineno: To get line number in code from node
• col-offset is the UTF-8 byte offset of the first token that generated the node

AST helpers:

• parse: Parse source code into AST node (alternative is compile). ast.parse(source,filename='<unknown>',mode='exec')Or compile(source,filename,mode,ast.PyCF_ONLY_AST)
• literal_eval: safely evaluate python expression. ast.literal_eval("{'name': 'Hello'}")
• walk: Recursively traverse to tree from root node. ast.walk(node)
• NodeVisitor class: It traverses to AST then call visit for each node. Node visitor class has some important methods: visit(node)which visit a node and generic_visit(node) which Visit all children of a node
• NodeTransformer class: Visit each node and transform it other form. In this case no need to use NodeVisitor class.
• There are a lot of other functions that AST provides, above mentioned classes/functions are majorly used.

Popular AST Nodes:

• Literals: Num(n), Str(s), List(elts, ctx), Tuple(elts, ctx), Set(elts), Dict(keys, value) etc…
• Variables: Name(id, ctx), Del, Store, Load and starred(value, ctx)
• Expressions: Expr(value), BinOP(left, op, right), Compare(left, ops, comparators), Call(func, args, keywords, starargs, kwargs), Attribute(value, attr, ctx), IfExp(test, body, orelse) etc…
• Subscripting: Subscript(value, slice, ctx), Index(value), Slice(lower, upper, step) etc…
• Statements: Assign(targets, value), Print(dest, values, nl), Delete(targets), Assert(test, msg), Pass etc…
• Imports: Import(names), ImportFrom(module, names, level) and alias(name, asname)
• Control flow: If(test, body, orelse), For(target, iter, body, orelse), While(test, body, orelse), Break, Continue, TryExcept(body, handlers, orelse), With(items, body) etc…
• Functions: FunctionDef(name, args, body, decorator_list, returns), Lambda(args, body), Return(value), Yield(value) etc…
• Classes: ClassDef(name, bases, keywords, starargs, kwargs, body, decorator_list)

AST code explanation with example:

Importing ast python library

To visit each node, at top NodeVisitor class is used. After that, the visit function definition for each node is defined.

1. visit_Import: For import any python library. import pandas
2. visit_ImportFrom: ImportFrom node. from datetime import datetime
3. visit_Assign: Definition for assignment. a = 5
4. visit_BinOp: visit definition for binary operation. a=3*6
5. visit_Expr: visit definition for expression. 3*7+5
6. visit_Num: visit definition for numbers
7. visit_Name: Definition for assigning values to a variable. a=5 variable “a” holding the value 5.
8. visit_Str: For string.

content is code text which you want to parse. Call NodeVisitor class to start traverse through the whole tree.

Sample output:

AST output

NodeTransformer class: You can convert the type of particular node to another node type using the NodeTransform class of AST. copy_location is used to copy source node location to the new node.

Abstract Syntax Tree can be used for other programming languages as well, but in this article, I focus more on Python AST working.

Possible use-case for Python AST:

Now, most of you have a question in your mind that where you can use this library and how it can be useful to you. So, I am going to explain to you one scenario:
Suppose you go to some online learning platform, where you want to learn how to write a program. This learning portal should have the capability to validate your code with the actual expected code. So, how you can do this sort of comparisons or validations. Some people might be thinking that do direct string comparisons why to stuck with some python library. I will be explaining you a major drawback if you do string comparisons for this sort of thing.
For e.g., the actual code is a = b+3 and the user’s code is a = 3+b, even in this simple scenario your string comparison will fail. How to do using ast library,
a = b + 3 or a = 3+b , both have same node type i.e. BinOp, you can validate variable “a” value and its node type.

For each line of code, create AST node then compare value, node type and other parameters as well like operator, operand, function name, class name, index, etc… if required.
Suppose you import some libraries in your code, easily you can find out details about your library like name, imported from and alias name using astlibrary.

Wrap up:

By now, hopefully, you understood

1. How to build an abstract syntax tree for your python code.
2. Types of nodes that can be used in AST.
3. Understanding of NodeVisitor and NodeTransform AST classes.

Source:

Official python documentation