CODEWITHSUNDEEP
pythonabstract-syntax-tree
Basssprosse
Basssprosse

Reputation: 466

Python's AST-Function

i want to analyse the following part with python's ast-functions:

extra_host_conf.setdefault('alias', []).extend( [(u'alias1', ['host1']), (u'alias2', ['host2']), (u'alias3', ['host3'])])

I need the list from the extend()-method-parameter. But it is important to check the string-parameter from the setdefault()-method: 'alias'! There could be also another string-parameter as 'alias'.

I know how to get the string 'alias' and i know how to get the list. But i dont know how to check the string and afterwards getting the list:(

That's how i get the string 'alias':

import ast

class MyVisitor(ast.NodeVisitor):
    def visit_Call(self, node):
        if node.func.attr == "setdefault":
            print ast.literal_eval(node.args[0])
        self.generic_visit(node)

And that is how i get the list:

import ast
    class MyVisitor(ast.NodeVisitor):
        def visit_Call(self, node):
            if node.func.attr == "extend":
                print ast.literal_eval(node.args[0])
            self.generic_visit(node)

But how to put both together? Like: "If the string is 'alias' please give me the list."

I wrote some methods to get the types of the different parts. I hope this can help:

ast.literal: [(u'alias1', ['host1']), (u'alias2', ['host2']), (u'alias3', ['host3'])]
Type: Call

Attribute: extend
Type: Attribute

ast.literal: alias
Type: Call

Attribute: setdefault
Type: Attribute

Name: extra_host_conf
Type: Name

String: alias
Type: Str

Upvotes: 4

Views: 3210

Answers (1)

Manthan Thakar
Manthan Thakar

Reputation: 209

If you dump the ast for the code snippet you've provided it looks like this:

Call(func=Attribute(value=Call(func=Attribute(value=Name(id='extra_host_conf', ctx=Load()), attr='setdefault', ctx=Load()), args=[Str(s='alias'), List(elts=[], ctx=Load())], keywords=[], starargs=None, kwargs=None), attr='extend', ctx=Load()), args=[List(elts=[Tuple(elts=[Str(s=u'alias1'), List(elts=[Str(s='host1')], ctx=Load())], ctx=Load()), Tuple(elts=[Str(s=u'alias2'), List(elts=[Str(s='host2')], ctx=Load())], ctx=Load()), Tuple(elts=[Str(s=u'alias3'), List(elts=[Str(s='host3')], ctx=Load())], ctx=Load())], ctx=Load())], keywords=[], starargs=None, kwargs=None)

Using this information you can check the appropriate condition to achieve the results you want.

Here's the code snippet that works for your usecase

import ast


class Visitor(ast.NodeVisitor):

    def visit_Call(self, node):
        if isinstance(node.func, ast.Attribute):
            func = node.func
            if func.value.func.attr == 'setdefault' and isinstance(func.value.args[0], ast.Str):
                # Now print the argument list associated with extend call
                if func.value.args[0].s == 'alias' and func.attr == 'extend':                 
                    print ast.dump(node.args[0])


    def run(self, code):
        n = ast.parse(code)
        self.visit(n)

if __name__ == '__main__':
    v = Visitor()
    v.run("extra_host_conf.setdefault('alias', []).extend( [(u'alias1', ['host1']), (u'alias2', ['host2']), (u'alias3', ['host3'])])")

Hope this helps!

Upvotes: 2

Related Questions