diff --git a/macropy/core/macros.py b/macropy/core/macros.py index 259356b4..67c5c529 100644 --- a/macropy/core/macros.py +++ b/macropy/core/macros.py @@ -178,6 +178,7 @@ def detect_macro(self, in_tree): in_tree.body, call_args, {'target': wi.optional_vars}, name) + in_tree.body = new_tree if new_tree: if isinstance(new_tree, ast.expr): diff --git a/macropy/core/test/macros/basic_block.py b/macropy/core/test/macros/basic_block.py index bb437757..1beb4db6 100644 --- a/macropy/core/test/macros/basic_block.py +++ b/macropy/core/test/macros/basic_block.py @@ -1,7 +1,19 @@ -from macropy.core.test.macros.basic_block_macro import macros, my_macro +from macropy.core.test.macros.basic_block_macro import macros, my_macro, my_nested_outer, my_nested_inner def run(): x = 10 with my_macro as y: x = x + 1 - return x \ No newline at end of file + + z = 10 + with my_nested_outer: + with my_nested_inner: + z = z + 1 # should get replaced by z = z * 2 + assert z == 80 + + z = 10 + with my_nested_inner, my_nested_outer: + z = z + 1 + assert z == 80 + + return x diff --git a/macropy/core/test/macros/basic_block_macro.py b/macropy/core/test/macros/basic_block_macro.py index ff42b534..8395fe48 100644 --- a/macropy/core/test/macros/basic_block_macro.py +++ b/macropy/core/test/macros/basic_block_macro.py @@ -8,3 +8,12 @@ def my_macro(tree, target, **kw): assert macropy.core.unparse(target) == "y" assert macropy.core.unparse(tree).strip() == "x = (x + 1)", macropy.core.unparse(tree) return tree * 3 + +@macros.block +def my_nested_outer(tree, **kw): + assert macropy.core.unparse(tree).strip() == "z = (z * 2)", macropy.core.unparse(tree) + return tree * 3 + +@macros.block +def my_nested_inner(tree, **kw): # important: generate new tree from scratch for this test + return [macropy.core.parse_stmt("z = z * 2")]