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

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

You can use ast.dump to get more information about the the nodes and AST 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))))])"

With above information you can change the visiting order of the node's children which allows you to generate postfix or prefix expression. In order to generate postfix expression change visit_BinOp, visit_BoolOp and visit_Call so that they visit arguments before visiting the operator/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 above changes you get following 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*********************************

['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 prefix expression instead just swap the order around so operator/function is visited 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']

Recommended topics

Hot tags