unify visitors for if and loop statements
diff --git a/mccabe.py b/mccabe.py
index eceafe6..6291320 100644
--- a/mccabe.py
+++ b/mccabe.py
@@ -162,32 +162,29 @@
def visitLoop(self, node):
name = "Loop %d" % node.lineno
-
- if self.graph is None:
- # global loop
- self.graph = PathGraph(name, name, node.lineno)
- pathnode = PathNode(name)
- self.tail = pathnode
- self.dispatch_list(node.body)
- self.graphs["%s%s" % (self.classname, name)] = self.graph
- self.reset()
- else:
- pathnode = self.appendPathNode(name)
- self.tail = pathnode
- self.dispatch_list(node.body)
- bottom = PathNode("", look='point')
- self.graph.connect(self.tail, bottom)
- self.graph.connect(pathnode, bottom)
- self.tail = bottom
-
- # TODO: else clause in node.orelse
+ self._inner1(node, name)
visitFor = visitWhile = visitLoop
def visitIf(self, node):
name = "If %d" % node.lineno
- pathnode = self.appendPathNode(name)
+ self._inner1(node, name)
+
+ def _inner1(self, node, name):
+ if self.graph is None:
+ # global loop
+ self.graph = PathGraph(name, name, node.lineno)
+ pathnode = PathNode(name)
+ self._inner2(node, pathnode)
+ self.graphs["%s%s" % (self.classname, name)] = self.graph
+ self.reset()
+ else:
+ pathnode = self.appendPathNode(name)
+ self._inner2(node, pathnode)
+
+ def _inner2(self, node, pathnode):
loose_ends = []
+ self.tail = pathnode
self.dispatch_list(node.body)
loose_ends.append(self.tail)
if node.orelse:
diff --git a/test_mccabe.py b/test_mccabe.py
index 1b9fc9c..6887c7c 100644
--- a/test_mccabe.py
+++ b/test_mccabe.py
@@ -45,6 +45,15 @@
"""
+for_else = """\
+def f(mylist):
+ for i in mylist:
+ print(i)
+ else:
+ print(None)
+"""
+
+
recursive = """\
def f(n):
if n > 4:
@@ -106,7 +115,12 @@
def test_for_loop_snippet(self):
complexity = get_complexity_number(for_loop, self.strio)
- # The for loop doesn't add an execution path, yet is counted.
+ # The for loop doesn't add an execution path
+ self.assertEqual(complexity, 1)
+
+ def test_for_else_snippet(self):
+ complexity = get_complexity_number(for_else, self.strio)
+ # The for loop doesn't add an execution path, but its `else` does
self.assertEqual(complexity, 2)
def test_recursive_snippet(self):