From 7289c18a42cc480b34838b7ff595c01b0a5baae9 Mon Sep 17 00:00:00 2001 From: Sebastien Bourdeauducq Date: Sat, 27 Sep 2014 11:57:47 +0800 Subject: [PATCH] transforms/unroll_loops: count statements in try blocks --- artiq/transforms/unroll_loops.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/artiq/transforms/unroll_loops.py b/artiq/transforms/unroll_loops.py index b362bd907..3c0f1de5d 100644 --- a/artiq/transforms/unroll_loops.py +++ b/artiq/transforms/unroll_loops.py @@ -4,12 +4,19 @@ from artiq.transforms.tools import eval_ast, value_to_ast 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