How to convert from infix to postfix / prefix using python AST module?

Question

I am trying to convert python math expressions to postfix notation using python AST module. Here is what I got so far:

import parser import ast from math import sin, cos, tan formulas = [ "1+2", "1+2*3", "1/2", "(1+2)*3", "sin(x)*x**2", "cos(x)", "True and False", "sin(w*time)" ] class v(ast.NodeVisitor): def __init__(self): self.tokens = [] def f_continue(self, node): super(v, self).generic_visit(node) def visit_Add(self, node): self.tokens.append('+') self.f_continue(node) def visit_And(self, node): self.tokens.append('&&') self.f_continue(node) def visit_BinOp(self, node): # print('visit_BinOp') # for child in ast.iter_fields(node): # print(' child %s ' % str(child)) self.f_continue(node) def visit_BoolOp(self, node): # print('visit_BoolOp') self.f_continue(node) def visit_Call(self, node): # print('visit_Call') self.f_continue(node) def visit_Div(self, node): self.tokens.append('/') self.f_continue(node) def visit_Expr(self, node): # print('visit_Expr') self.f_continue(node) def visit_Import(self, stmt_import): for alias in stmt_import.names: print('import name "%s"' % alias.name) print('import object %s' % alias) self.f_continue(stmt_import) def visit_Load(self, node): # print('visit_Load') self.f_continue(node) def visit_Module(self, node): # print('visit_Module') self.f_continue(node) def visit_Mult(self, node): self.tokens.append('*') self.f_continue(node) def visit_Name(self, node): self.tokens.append(node.id) self.f_continue(node) def visit_NameConstant(self, node): self.tokens.append(node.value) self.f_continue(node) def visit_Num(self, node): self.tokens.append(node.n) self.f_continue(node) def visit_Pow(self, node): self.tokens.append('pow') self.f_continue(node) for index, f in enumerate(formulas): print('{} - {:*^76}'.format(index, f)) visitor = v() visitor.visit(ast.parse(f)) print(visitor.tokens) print() # 0 - ************************************1+2************************************* # [1, '+', 2] # 1 - ***********************************1+2*3************************************ # [1, '+', 2, '*', 3] # 2 - ************************************1/2************************************* # [1, '/', 2] # 3 - **********************************(1+2)*3*********************************** # [1, '+', 2, '*', 3] # 4 - ********************************sin(x)*x**2********************************* # ['sin', 'x', '*', 'x', 'pow', 2] # 5 - ***********************************cos(x)*********************************** # ['cos', 'x'] # 6 - *******************************True and False******************************* # ['&&', True, False] # 7 - ********************************sin(w*time)********************************* # ['sin', 'w', '*', 'time'] 

I am trying to understand how to convert complex mathematical expressions infix into postfix expressions for sending to the swig c shell, to do this, I am trying to use the AST module.

Can anyone advise here?

Answer 1

You can use ast.dump to get more information about AST nodes and structure:

 >>> import ast >>> node = ast.parse("sin(x)*x**2") >>> ast.dump(node) "Module(body=[Expr(value=BinOp(left=Call(func=Name(id='sin', ctx=Load()), args=[Name(id='x', ctx=Load())], keywords=[]), op=Mult(), right=BinOp(left=Name(id='x', ctx=Load()), op=Pow(), right=Num(n=2))))])" 

Using the above information, you can change the order of visits to the children of a node, which allows you to generate a postfix or prefix expression. To generate the postfix expression change visit_BinOp , visit_BoolOp and visit_Call so that they visit the arguments before visiting the statement / function:

 def visit_BinOp(self, node): self.visit(node.left) self.visit(node.right) self.visit(node.op) def visit_BoolOp(self, node): for val in node.values: self.visit(val) self.visit(node.op) def visit_Call(self, node): for arg in node.args: self.visit(arg) self.visit(node.func) 

With the above changes, you get the following result:

 0 - ************************************1+2************************************* [1, 2, '+'] 1 - ***********************************1+2*3************************************ [1, 2, 3, '*', '+'] 2 - ************************************1/2************************************* [1, 2, '/'] 3 - **********************************(1+2)*3*********************************** [1, 2, '+', 3, '*'] 4 - ********************************sin(x)*x**2********************************* ['x', 'sin', 'x', 2, 'pow', '*'] 5 - ***********************************cos(x)*********************************** ['x', 'cos'] 6 - *******************************True and False******************************* [True, False, '&&'] 7 - ********************************sin(w*time)********************************* ['w', 'time', '*', 'sin'] 

If you want to use a prefix expression instead of just changing the order, since the operator / function is searched first:

 def visit_BinOp(self, node): self.visit(node.op) self.visit(node.left) self.visit(node.right) def visit_BoolOp(self, node): self.visit(node.op) for val in node.values: self.visit(val) def visit_Call(self, node): self.visit(node.func) for arg in node.args: self.visit(arg) 

Output:

 0 - ************************************1+2************************************* ['+', 1, 2] 1 - ***********************************1+2*3************************************ ['+', 1, '*', 2, 3] 2 - ************************************1/2************************************* ['/', 1, 2] 3 - **********************************(1+2)*3*********************************** ['*', '+', 1, 2, 3] 4 - ********************************sin(x)*x**2********************************* ['*', 'sin', 'x', 'pow', 'x', 2] 5 - ***********************************cos(x)*********************************** ['cos', 'x'] 6 - *******************************True and False******************************* ['&&', True, False] 7 - ********************************sin(w*time)********************************* ['sin', '*', 'w', 'time']