CODEWITHSUNDEEP
pythonabstract-syntax-tree
Harry Adams
Harry Adams

Reputation: 487

Python ast - getting function parameters and processing them

I'm trying to use the ast module in Python to parse input code, but am struggling with a lot of the syntax of how to do so. For instance, I have the following code as a testing environment:

import ast


class NodeVisitor(ast.NodeVisitor):
    def visit_Call(self, node):
        for each in node.args:
            print(ast.literal_eval(each))
        self.generic_visit(node)


line = "circuit = QubitCircuit(3, True)"
tree = ast.parse(line)

print("VISITOR")
visitor = NodeVisitor()
visitor.visit(tree)

Output:

VISITOR
3
True

In this instance, and please correct me if I'm wrong, the visit_Call will be used if it's a function call? So I can get each argument, however there's no guarantee it will work like this as there are different arguments available to be provided. I understand that node.args is providing my arguments, but I'm not sure how to do things with them?

I guess what I'm asking is how do I check what the arguments are and do different things with them? I'd like to check, perhaps, that the first argument is an Int, and if so, run processInt(parameter) as an example.

Upvotes: 0

Views: 2896

Answers (1)

Blckknght
Blckknght

Reputation: 104712

The value each in your loop in the method will be assigned to the AST node for each of the arguments in each function call you visit. There are lots of different types of AST nodes, so by checking which kind you have, you may be able to learn things about the argument being passed in.

Note however that the AST is about syntax, not values. So if the function call was foo(bar), it's just going to tell you that the argument is a variable named bar, not what the value of that variable is (which it does not know). If the function call was foo(bar(baz)), it's going to show you that the argument is another function call. If you only need to handle calls with literals as their arguments, then you're probably going to be OK, you'll just look instances of AST.Num and similar.

If you want to check if the first argument is a number and process it if it is, you can do something like:

def visit_Call(self, node):
    first_arg = node.args[0]
    if isinstance(first_arg, ast.Num):
        processInt(first_arg.n)
    else:
        pass  # Do you want to do something on a bad argument? Raise an exception maybe?

Upvotes: 2

Related Questions