pythonpylintastroid

How to get number of conditions in if statement using Astroid?


I am working on Pylint project using Astroid package for Python. I would like to know if there is a general and convenient way to retrieve the number of conditions branches of an if-statement using Astroid

For instance, for the statement:

if (branch1 and branch2): pass

I would like the returned value 2.

I found a naïve way to do it using astroid:

 def _has_one_condition(node):
        return isinstance(node, astroid.Compare)
   
 def _has_2_conditions(node):
        
        return (
                isinstance(node, astroid.BoolOp)
                and len(node.values) == 2
                and isinstance(node.values[0], astroid.Compare)
                and isinstance(node.values[1], astroid.Compare)
        )
    

I am looking for a more general way to do this with astroid


Solution

  • This is a way to count the number of branches of an if statement. Make sure that the argument being passed as node is the .test attribute of the if astroid node.

    def num_of_condition_branches(node):
        if isinstance(node, astroid.BoolOp):
            return _num_of_condition_branches_helper(node,0)
        return 1
    
    
    def _num_of_condition_branches_helper(node,branches):
        for child in node.get_children():
            if not isinstance(child, astroid.BoolOp):
                branches+=1
            else:
                branches=_num_of_condition_branches_helper(child,branches)
        return branches