Replace an expression containing '/' with div function call in Python

Question

I am working on Python code to replace an expression containing '/' with actual function call. For eg: '(n/7-7) +(n/3+3)' should become '(div(n,7)-7 + ( div(n,3)+3)'. Please note only '/' operand needs to be replaced.

I am using ast.NodeVisitor for the same.

class visitor(ast.NodeVisitor):
    def visit_Expression(self, node): return self.visit(node.body)
    def visit_Name(self, node): return node.id
    def visit_Constant(self, node): return str(node.value)
    def visit_UnaryOp(self, node): return f'{self.visit(node.op)}({self.visit(node.operand)})'
    def visit_UAdd(self, node): return '+'
    def visit_USub(self, node): return '-'
    def visit_BinOp(self, node): 
        if isinstance(node.op,ast.Div):
             return f'{self.visit(node.op)}({self.visit(node.left)},{self.visit(node.right)})'
        else:
            return f'{self.visit(node.left)} {self.visit(node.op)} {self.visit(node.right)}'
    def visit_Add(self, node): return '+'
    def visit_Sub(self, node): return '-'
    def visit_Mult(self, node): return '*'
    def visit_Div(self, node): return 'div'
    def generic_visit(self, node): raise ValueError('Invalid Token')

def tokenize(source):
    return visitor().visit(ast.parse(source, mode='eval'))

I am calling this function as and get the output div(n,5)-1:

    expression = 'n/5-1'
    expression = tokenize(expression)

This however does not work for trigonometric functions like tan/radians For eg:

expression = 'tan(radians(top))'

Do I need to add visit_tan and visit_radians as well?

Answer 1

If you intend to continue with your current approach, I don't see a reason you need anything else than you current extension of ast.NodeVisitor.

You could just add the following method to your class:

class visitor(ast.NodeVisitor):
    # (...)
    def visit_Call(self, node): 
        return f'{self.visit(node.func)}({','.join(map(self.visit, node.args))})'

tokenize('tan(radians(n/7-7))')
# 'tan(radians(div(n,7) - 7))'

This will read any function call (like rad, tan...) and write them again as function call in the resulting text. If you want to restrict them to allow only a specific list of function names, you could put a check inside the method, and only dispatch to visit(node.func) if node.func.id is in your list, otherwise, dispatch to the generic_visit so it can throw an error.