transforms/unroll_loops: count statements in try blocks

sbourdeauducq · sbourdeauducq · commit 7289c18a42cc · 2014-09-27T11:57:47.000+08:00
diff --git a/artiq/transforms/unroll_loops.py b/artiq/transforms/unroll_loops.py
@@ -4,12 +4,19 @@
 
 
 def _count_stmts(node):
-    if isinstance(node, (ast.For, ast.While, ast.If)):
-        return 1 + _count_stmts(node.body) + _count_stmts(node.orelse)
+    if isinstance(node, list):
+        return sum(map(_count_stmts, node))
     elif isinstance(node, ast.With):
         return 1 + _count_stmts(node.body)
-    elif isinstance(node, list):
-        return sum(map(_count_stmts, node))
+    elif isinstance(node, (ast.For, ast.While, ast.If)):
+        return 1 + _count_stmts(node.body) + _count_stmts(node.orelse)
+    elif isinstance(node, ast.Try):
+        r = 1 + _count_stmts(node.body) \
+            + _count_stmts(node.orelse) \
+            + _count_stmts(node.finalbody)
+        for handler in node.handlers:
+            r += 1 + _count_stmts(handler.body)
+        return r
     else:
         return 1
 
