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

Question

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

An `if` statement accepts exactly one condition. This condition might consist of any kind of expression with any number of operators in any number of nesting levels with or without recursion. How exactly should that be counted? — Klaus D., Nov 03 '20 at 13:19
I edited the question - I wanted to write a function that returns the numbers of branches in a condition using the `astroid` package. — Or b, Jan 01 '21 at 23:39

score 0 · Accepted Answer · answered Jan 01 '21 at 23:35

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

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

1 Answers1