How to use Python's ast library to detect all mathematical operations and operands used within a function?

Viewed 107

I'm writing an application in Python (version : Python 3.8.8). My main goal is to implement a script in Python that detects different mathematical operations (addition, multplication, etc.) performed within a function. At the same time, I also need to know the operands used in each operation.

Let's take an example. We have a function called evaluate_expr(expr) to detect all operations and operands when we want to test for example an expression by doing this :

node = ast.parse("3+5", mode="eval")
evaluate_expr(node.body)

Below we have the definition of the function evaluate_expr(expr):

def evaluate_expr(expr):
    if isinstance(expr, ast.Expression):
       return evaluate_expr(expr.body)
    elif isinstance(expr, ast.Num):
       return expr.n
    elif isinstance(expr, ast.BinOp):
       op = expr.op
       left = evaluate_expr(expr.left)
       right = evaluate_expr(expr.right)
       print("left : ",left)
       print("right : ", right)
    if isinstance(op, ast.Add):
       return left + right
    if isinstance(op, ast.Sub):
       return left - right
    elif isinstance(op, ast.Mult):
        return left * right
    elif isinstance(op, ast.Div):
        return left / right
raise ValueError(f"Can't evaluate {expr}")

The output of print(evaluate_expr(node.body) is :

left : 3
right : 5
8

which is correct.

Now my difficulties is when I want to define a function performing one or more operations. How can I use ast library ? For example, if I define a function like this :

def foo(x, y=2):
  z = x*y + 3
  print("z is ", z)
  return z**2

Do you have an approach to propose in order to find in each row within the function all the operations calculated successively with operands ? For example, if we want to test foo(3,2) with ast library of Python, I would like to get in output this :

leftOperand : 3
rightOperand : 2
result :6 (as we are computing x*y)

then :

leftOperand : 6
rightOperand : 3
result : 9 (as we are computing x*y + 3, and we have x*y =  6, so we have 6+3=9)

then :

result : z is 9

then :

leftOperand : 9
rightOperand : 9
result : 81 (as we are computing z**2, 9*9)

Do you have any advice ?

0 Answers
Related