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
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