forked from M-Labs/artiq
1
0
Fork 0

transforms/unroll_loop: handle breaks in try and with blocks

This commit is contained in:
Sebastien Bourdeauducq 2014-09-26 19:20:56 +08:00
parent f332def3a2
commit 8852dd394e
1 changed files with 17 additions and 13 deletions

View File

@ -15,21 +15,25 @@ def _count_stmts(node):
def _loop_breakable(node): def _loop_breakable(node):
if isinstance(node, ast.Break): if isinstance(node, list):
return 1 return any(map(_loop_breakable, node))
elif isinstance(node, ast.Return): elif isinstance(node, (ast.Break, ast.Continue)):
return 2 return True
elif isinstance(node, list): elif isinstance(node, ast.With):
return max(map(_loop_breakable, node), default=0) return _loop_breakable(node.body)
elif isinstance(node, ast.If): elif isinstance(node, ast.If):
return max(_loop_breakable(node.body), _loop_breakable(node.orelse)) return _loop_breakable(node.body) or _loop_breakable(node.orelse)
elif isinstance(node, (ast.For, ast.While)): elif isinstance(node, ast.Try):
bb = _loop_breakable(node.body) if (_loop_breakable(node.body)
if bb == 1: or _loop_breakable(node.orelse)
bb = 0 or _loop_breakable(node.finalbody)):
return max(bb, _loop_breakable(node.orelse)) return True
for handler in node.handlers:
if _loop_breakable(handler.body):
return True
return False
else: else:
return 0 return False
class _LoopUnroller(ast.NodeTransformer): class _LoopUnroller(ast.NodeTransformer):