fix: fuse with statements in codegen

bolt/codegen.py

@@ -98,6 +98,32 @@ class CodegenStatement:
     lineno: Optional[int] = None
     children: List["CodegenStatement"] = field(default_factory=list)
 
+    def is_with_statement(self) -> bool:
+        """Check if the statement is a with statement."""
+        return (
+            self.code.startswith("with ")
+            and self.code.endswith(":")
+            and bool(self.children)
+        )
+
+    def fuse_with_statements(self) -> "CodegenStatement":
+        """Fuse nested with statements."""
+        if (
+            self.is_with_statement()
+            and len(self.children) == 1
+            and self.children[0].is_with_statement()
+        ):
+            outer = self.code[5:-1]
+            inner = self.children[0].code[5:-1]
+            return replace(
+                self.children[0],
+                code=f"with {outer}, {inner}:",
+                lineno=self.lineno or self.children[0].lineno,
+            ).fuse_with_statements()
+        return replace(
+            self, children=[child.fuse_with_statements() for child in self.children]
+        )
+
     def flatten(self, indent: str = "") -> Iterable[Tuple[str, Optional[int]]]:
         """Yield the indented statements with their associated line number."""
         yield f"{indent}{self.code}", self.lineno
@@ -138,7 +164,7 @@ def get_source(self) -> str:
         numbers2: List[int] = [1]
 
         for statement in header + self.statements:
-            for code, lineno in statement.flatten():
+            for code, lineno in statement.fuse_with_statements().flatten():
                 if lineno and numbers2[-1] != lineno:
                     numbers1.append(len(lines) + 1)
                     numbers2.append(lineno)

tests/snapshots/bolt__parse_103__1.txt

@@ -1,17 +1,13 @@
-_bolt_lineno = [1, 13], [1, 4]
+_bolt_lineno = [1, 9], [1, 4]
 _bolt_helper_interpolate_nbt = _bolt_runtime.helpers['interpolate_nbt']
 _bolt_helper_replace = _bolt_runtime.helpers['replace']
 _bolt_helper_children = _bolt_runtime.helpers['children']
 with _bolt_runtime.scope() as _bolt_var2:
     _bolt_var0 = '[{"text": "Hello", "bold": true}]'
     foo = _bolt_var0
-    with _bolt_runtime.push_nesting('execute:subcommand'):
-        with _bolt_runtime.push_nesting('execute:as:targets:subcommand', *_bolt_refs[15:16]):
-            with _bolt_runtime.push_nesting('execute:at:targets:subcommand', *_bolt_refs[13:14]):
-                with _bolt_runtime.push_nesting('execute:if:block:pos:block:subcommand', *_bolt_refs[10:12]):
-                    with _bolt_runtime.push_nesting('execute:run:subcommand'):
-                        _bolt_var1 = foo
-                        _bolt_var1 = _bolt_helper_interpolate_nbt(_bolt_var1, _bolt_refs[0])
+    with _bolt_runtime.push_nesting('execute:subcommand'), _bolt_runtime.push_nesting('execute:as:targets:subcommand', *_bolt_refs[15:16]), _bolt_runtime.push_nesting('execute:at:targets:subcommand', *_bolt_refs[13:14]), _bolt_runtime.push_nesting('execute:if:block:pos:block:subcommand', *_bolt_refs[10:12]), _bolt_runtime.push_nesting('execute:run:subcommand'):
+        _bolt_var1 = foo
+        _bolt_var1 = _bolt_helper_interpolate_nbt(_bolt_var1, _bolt_refs[0])
     _bolt_runtime.commands.append(_bolt_helper_replace(_bolt_refs[17], arguments=_bolt_helper_children([_bolt_helper_replace(_bolt_refs[16], arguments=_bolt_helper_children([*_bolt_refs[15:16], _bolt_helper_replace(_bolt_refs[14], arguments=_bolt_helper_children([*_bolt_refs[13:14], _bolt_helper_replace(_bolt_refs[12], arguments=_bolt_helper_children([*_bolt_refs[10:12], _bolt_helper_replace(_bolt_refs[9], arguments=_bolt_helper_children([_bolt_helper_replace(_bolt_refs[8], arguments=_bolt_helper_children([_bolt_refs[7], _bolt_helper_replace(_bolt_refs[6], data_tags=_bolt_helper_replace(_bolt_refs[5], entries=_bolt_helper_children([_bolt_helper_replace(_bolt_refs[4], value=_bolt_helper_replace(_bolt_refs[3], entries=_bolt_helper_children([_bolt_helper_replace(_bolt_refs[1], value=_bolt_var1), _bolt_refs[2]])))])))]))]))]))]))]))])))
 _bolt_var3 = _bolt_helper_replace(_bolt_refs[18], commands=_bolt_helper_children(_bolt_var2))
 ---

tests/snapshots/bolt__parse_130__1.txt

@@ -1,22 +1,18 @@
-_bolt_lineno = [1, 13], [1, 2]
+_bolt_lineno = [1, 9], [1, 2]
 _bolt_helper_operator_not = _bolt_runtime.helpers['operator_not']
 _bolt_helper_branch = _bolt_runtime.helpers['branch']
 _bolt_helper_children = _bolt_runtime.helpers['children']
 _bolt_helper_replace = _bolt_runtime.helpers['replace']
 with _bolt_runtime.scope() as _bolt_var2:
-    with _bolt_runtime.push_nesting('execute:subcommand'):
-        with _bolt_runtime.push_nesting('execute:if:score:target:targetObjective:matches:range:subcommand', *_bolt_refs[6:9]):
-            with _bolt_runtime.push_nesting('execute:if:entity:entities:subcommand', *_bolt_refs[4:5]):
-                with _bolt_runtime.push_nesting('execute:commands'):
-                    with _bolt_runtime.scope() as _bolt_var1:
-                        _bolt_var0 = True
-                        _bolt_var0_inverse = _bolt_helper_operator_not(_bolt_var0)
-                        with _bolt_helper_branch(_bolt_var0) as _bolt_condition:
-                            if _bolt_condition:
-                                _bolt_runtime.commands.extend(_bolt_refs[0].commands)
-                        with _bolt_helper_branch(_bolt_var0_inverse) as _bolt_condition:
-                            if _bolt_condition:
-                                _bolt_runtime.commands.extend(_bolt_refs[1].commands)
+    with _bolt_runtime.push_nesting('execute:subcommand'), _bolt_runtime.push_nesting('execute:if:score:target:targetObjective:matches:range:subcommand', *_bolt_refs[6:9]), _bolt_runtime.push_nesting('execute:if:entity:entities:subcommand', *_bolt_refs[4:5]), _bolt_runtime.push_nesting('execute:commands'), _bolt_runtime.scope() as _bolt_var1:
+        _bolt_var0 = True
+        _bolt_var0_inverse = _bolt_helper_operator_not(_bolt_var0)
+        with _bolt_helper_branch(_bolt_var0) as _bolt_condition:
+            if _bolt_condition:
+                _bolt_runtime.commands.extend(_bolt_refs[0].commands)
+        with _bolt_helper_branch(_bolt_var0_inverse) as _bolt_condition:
+            if _bolt_condition:
+                _bolt_runtime.commands.extend(_bolt_refs[1].commands)
     _bolt_runtime.commands.append(_bolt_helper_replace(_bolt_refs[10], arguments=_bolt_helper_children([_bolt_helper_replace(_bolt_refs[9], arguments=_bolt_helper_children([*_bolt_refs[6:9], _bolt_helper_replace(_bolt_refs[5], arguments=_bolt_helper_children([*_bolt_refs[4:5], _bolt_helper_replace(_bolt_refs[3], arguments=_bolt_helper_children([_bolt_helper_replace(_bolt_refs[2], commands=_bolt_helper_children(_bolt_var1))]))]))]))])))
 _bolt_var3 = _bolt_helper_replace(_bolt_refs[11], commands=_bolt_helper_children(_bolt_var2))
 ---

tests/snapshots/bolt__parse_139__1.txt

@@ -1,4 +1,4 @@
-_bolt_lineno = [1, 15, 20, 28, 40], [1, 2, 3, 4, 6]
+_bolt_lineno = [1, 15, 19, 27, 38], [1, 2, 3, 4, 6]
 _bolt_helper_get_attribute_handler = _bolt_runtime.helpers['get_attribute_handler']
 _bolt_helper_interpolate_resource_location = _bolt_runtime.helpers['interpolate_resource_location']
 _bolt_helper_children = _bolt_runtime.helpers['children']
@@ -15,37 +15,34 @@ with _bolt_runtime.scope() as _bolt_var10:
         _bolt_var2 = node
         _bolt_var2 = _bolt_helper_get_attribute_handler(_bolt_var2)["parent"]
         _bolt_var2 = _bolt_helper_interpolate_resource_location(_bolt_var2, _bolt_refs[0])
-        with _bolt_runtime.push_nesting('append:function:name:commands', *_bolt_helper_children([_bolt_var2])):
-            with _bolt_runtime.scope() as _bolt_var9:
-                _bolt_var3 = node
-                _bolt_var3 = _bolt_helper_get_attribute_handler(_bolt_var3)["partition"]
-                _bolt_var4 = 5
-                _bolt_var3 = _bolt_var3(_bolt_var4)
-                _bolt_var3_inverse = _bolt_helper_operator_not(_bolt_var3)
-                with _bolt_helper_branch(_bolt_var3) as _bolt_condition:
-                    if _bolt_condition:
-                        with _bolt_runtime.push_nesting('execute:subcommand'):
-                            _bolt_var5 = node
-                            _bolt_var5 = _bolt_helper_get_attribute_handler(_bolt_var5)["range"]
-                            _bolt_var5 = _bolt_helper_interpolate_range(_bolt_var5, _bolt_refs[1])
-                            with _bolt_runtime.push_nesting('execute:if:score:target:targetObjective:matches:range:subcommand', *_bolt_helper_children([_bolt_refs[2], _bolt_refs[3], _bolt_var5])):
-                                with _bolt_runtime.push_nesting('execute:run:subcommand'):
-                                    _bolt_var6 = node
-                                    _bolt_var6 = _bolt_helper_get_attribute_handler(_bolt_var6)["children"]
-                                    _bolt_var6 = _bolt_helper_interpolate_resource_location(_bolt_var6, _bolt_refs[4])
-                        _bolt_runtime.commands.append(_bolt_helper_replace(_bolt_refs[8], arguments=_bolt_helper_children([_bolt_helper_replace(_bolt_refs[7], arguments=_bolt_helper_children([*_bolt_helper_children([_bolt_refs[2], _bolt_refs[3], _bolt_var5]), _bolt_helper_replace(_bolt_refs[6], arguments=_bolt_helper_children([_bolt_helper_replace(_bolt_refs[5], arguments=_bolt_helper_children([_bolt_var6]))]))]))])))
-                with _bolt_helper_branch(_bolt_var3_inverse) as _bolt_condition:
-                    if _bolt_condition:
-                        with _bolt_runtime.push_nesting('execute:subcommand'):
-                            _bolt_var7 = node
-                            _bolt_var7 = _bolt_helper_get_attribute_handler(_bolt_var7)["range"]
-                            _bolt_var7 = _bolt_helper_interpolate_range(_bolt_var7, _bolt_refs[9])
-                            with _bolt_runtime.push_nesting('execute:if:score:target:targetObjective:matches:range:subcommand', *_bolt_helper_children([_bolt_refs[10], _bolt_refs[11], _bolt_var7])):
-                                with _bolt_runtime.push_nesting('execute:run:subcommand'):
-                                    _bolt_var8 = node
-                                    _bolt_var8 = _bolt_helper_get_attribute_handler(_bolt_var8)["value"]
-                                    _bolt_var8 = _bolt_helper_interpolate_message(_bolt_var8, _bolt_refs[12])
-                        _bolt_runtime.commands.append(_bolt_helper_replace(_bolt_refs[16], arguments=_bolt_helper_children([_bolt_helper_replace(_bolt_refs[15], arguments=_bolt_helper_children([*_bolt_helper_children([_bolt_refs[10], _bolt_refs[11], _bolt_var7]), _bolt_helper_replace(_bolt_refs[14], arguments=_bolt_helper_children([_bolt_helper_replace(_bolt_refs[13], arguments=_bolt_helper_children([_bolt_var8]))]))]))])))
+        with _bolt_runtime.push_nesting('append:function:name:commands', *_bolt_helper_children([_bolt_var2])), _bolt_runtime.scope() as _bolt_var9:
+            _bolt_var3 = node
+            _bolt_var3 = _bolt_helper_get_attribute_handler(_bolt_var3)["partition"]
+            _bolt_var4 = 5
+            _bolt_var3 = _bolt_var3(_bolt_var4)
+            _bolt_var3_inverse = _bolt_helper_operator_not(_bolt_var3)
+            with _bolt_helper_branch(_bolt_var3) as _bolt_condition:
+                if _bolt_condition:
+                    with _bolt_runtime.push_nesting('execute:subcommand'):
+                        _bolt_var5 = node
+                        _bolt_var5 = _bolt_helper_get_attribute_handler(_bolt_var5)["range"]
+                        _bolt_var5 = _bolt_helper_interpolate_range(_bolt_var5, _bolt_refs[1])
+                        with _bolt_runtime.push_nesting('execute:if:score:target:targetObjective:matches:range:subcommand', *_bolt_helper_children([_bolt_refs[2], _bolt_refs[3], _bolt_var5])), _bolt_runtime.push_nesting('execute:run:subcommand'):
+                            _bolt_var6 = node
+                            _bolt_var6 = _bolt_helper_get_attribute_handler(_bolt_var6)["children"]
+                            _bolt_var6 = _bolt_helper_interpolate_resource_location(_bolt_var6, _bolt_refs[4])
+                    _bolt_runtime.commands.append(_bolt_helper_replace(_bolt_refs[8], arguments=_bolt_helper_children([_bolt_helper_replace(_bolt_refs[7], arguments=_bolt_helper_children([*_bolt_helper_children([_bolt_refs[2], _bolt_refs[3], _bolt_var5]), _bolt_helper_replace(_bolt_refs[6], arguments=_bolt_helper_children([_bolt_helper_replace(_bolt_refs[5], arguments=_bolt_helper_children([_bolt_var6]))]))]))])))
+            with _bolt_helper_branch(_bolt_var3_inverse) as _bolt_condition:
+                if _bolt_condition:
+                    with _bolt_runtime.push_nesting('execute:subcommand'):
+                        _bolt_var7 = node
+                        _bolt_var7 = _bolt_helper_get_attribute_handler(_bolt_var7)["range"]
+                        _bolt_var7 = _bolt_helper_interpolate_range(_bolt_var7, _bolt_refs[9])
+                        with _bolt_runtime.push_nesting('execute:if:score:target:targetObjective:matches:range:subcommand', *_bolt_helper_children([_bolt_refs[10], _bolt_refs[11], _bolt_var7])), _bolt_runtime.push_nesting('execute:run:subcommand'):
+                            _bolt_var8 = node
+                            _bolt_var8 = _bolt_helper_get_attribute_handler(_bolt_var8)["value"]
+                            _bolt_var8 = _bolt_helper_interpolate_message(_bolt_var8, _bolt_refs[12])
+                    _bolt_runtime.commands.append(_bolt_helper_replace(_bolt_refs[16], arguments=_bolt_helper_children([_bolt_helper_replace(_bolt_refs[15], arguments=_bolt_helper_children([*_bolt_helper_children([_bolt_refs[10], _bolt_refs[11], _bolt_var7]), _bolt_helper_replace(_bolt_refs[14], arguments=_bolt_helper_children([_bolt_helper_replace(_bolt_refs[13], arguments=_bolt_helper_children([_bolt_var8]))]))]))])))
         _bolt_runtime.commands.append(_bolt_helper_replace(_bolt_refs[18], arguments=_bolt_helper_children([*_bolt_helper_children([_bolt_var2]), _bolt_helper_replace(_bolt_refs[17], commands=_bolt_helper_children(_bolt_var9))])))
 _bolt_var11 = _bolt_helper_replace(_bolt_refs[19], commands=_bolt_helper_children(_bolt_var10))
 ---

tests/snapshots/bolt__parse_229__1.txt

@@ -1,14 +1,12 @@
-_bolt_lineno = [1, 11], [1, 3]
+_bolt_lineno = [1, 9], [1, 3]
 _bolt_helper_children = _bolt_runtime.helpers['children']
 _bolt_helper_macro_call = _bolt_runtime.helpers['macro_call']
 _bolt_helper_replace = _bolt_runtime.helpers['replace']
 with _bolt_runtime.scope() as _bolt_var1:
     def _bolt_macro0():
         pass
-    with _bolt_runtime.push_nesting('execute:subcommand'):
-        with _bolt_runtime.push_nesting('execute:as:targets:subcommand', *_bolt_refs[2:3]):
-            with _bolt_runtime.push_nesting('execute:run:subcommand'):
-                _bolt_var0 = _bolt_helper_macro_call(_bolt_runtime, _bolt_macro0, _bolt_refs[0])
+    with _bolt_runtime.push_nesting('execute:subcommand'), _bolt_runtime.push_nesting('execute:as:targets:subcommand', *_bolt_refs[2:3]), _bolt_runtime.push_nesting('execute:run:subcommand'):
+        _bolt_var0 = _bolt_helper_macro_call(_bolt_runtime, _bolt_macro0, _bolt_refs[0])
     _bolt_runtime.commands.append(_bolt_helper_replace(_bolt_refs[4], arguments=_bolt_helper_children([_bolt_helper_replace(_bolt_refs[3], arguments=_bolt_helper_children([*_bolt_refs[2:3], _bolt_helper_replace(_bolt_refs[1], arguments=_bolt_helper_children([*_bolt_var0]))]))])))
 _bolt_var2 = _bolt_helper_replace(_bolt_refs[5], commands=_bolt_helper_children(_bolt_var1))
 ---

tests/snapshots/bolt__parse_231__1.txt

@@ -1,13 +1,12 @@
-_bolt_lineno = [1, 10], [1, 3]
+_bolt_lineno = [1, 9], [1, 3]
 _bolt_helper_children = _bolt_runtime.helpers['children']
 _bolt_helper_macro_call = _bolt_runtime.helpers['macro_call']
 _bolt_helper_replace = _bolt_runtime.helpers['replace']
 with _bolt_runtime.scope() as _bolt_var1:
     def _bolt_macro0():
         pass
-    with _bolt_runtime.push_nesting('execute:subcommand'):
-        with _bolt_runtime.push_nesting('execute:as:targets:subcommand', *_bolt_refs[1:2]):
-            _bolt_var0 = _bolt_helper_macro_call(_bolt_runtime, _bolt_macro0, _bolt_refs[0])
+    with _bolt_runtime.push_nesting('execute:subcommand'), _bolt_runtime.push_nesting('execute:as:targets:subcommand', *_bolt_refs[1:2]):
+        _bolt_var0 = _bolt_helper_macro_call(_bolt_runtime, _bolt_macro0, _bolt_refs[0])
     _bolt_runtime.commands.append(_bolt_helper_replace(_bolt_refs[3], arguments=_bolt_helper_children([_bolt_helper_replace(_bolt_refs[2], arguments=_bolt_helper_children([*_bolt_refs[1:2], *_bolt_var0]))])))
 _bolt_var2 = _bolt_helper_replace(_bolt_refs[4], commands=_bolt_helper_children(_bolt_var1))
 ---

tests/snapshots/bolt__parse_355__1.txt

@@ -1,13 +1,12 @@
-_bolt_lineno = [1, 10], [1, 3]
+_bolt_lineno = [1, 9], [1, 3]
 _bolt_helper_children = _bolt_runtime.helpers['children']
 _bolt_helper_macro_call = _bolt_runtime.helpers['macro_call']
 _bolt_helper_replace = _bolt_runtime.helpers['replace']
 with _bolt_runtime.scope() as _bolt_var1:
     def _bolt_macro0(thing):
         pass
-    with _bolt_runtime.push_nesting('execute:subcommand'):
-        with _bolt_runtime.push_nesting('execute:as:targets:subcommand', *_bolt_refs[1:2]):
-            _bolt_var0 = _bolt_helper_macro_call(_bolt_runtime, _bolt_macro0, _bolt_refs[0])
+    with _bolt_runtime.push_nesting('execute:subcommand'), _bolt_runtime.push_nesting('execute:as:targets:subcommand', *_bolt_refs[1:2]):
+        _bolt_var0 = _bolt_helper_macro_call(_bolt_runtime, _bolt_macro0, _bolt_refs[0])
     _bolt_runtime.commands.append(_bolt_helper_replace(_bolt_refs[3], arguments=_bolt_helper_children([_bolt_helper_replace(_bolt_refs[2], arguments=_bolt_helper_children([*_bolt_refs[1:2], *_bolt_var0]))])))
 _bolt_var2 = _bolt_helper_replace(_bolt_refs[4], commands=_bolt_helper_children(_bolt_var1))
 ---

tests/snapshots/bolt__parse_42__1.txt

@@ -1,24 +1,21 @@
-_bolt_lineno = [1, 12, 18], [1, 2, 4]
+_bolt_lineno = [1, 9, 15], [1, 2, 4]
 _bolt_helper_branch = _bolt_runtime.helpers['branch']
 _bolt_helper_children = _bolt_runtime.helpers['children']
 _bolt_helper_replace = _bolt_runtime.helpers['replace']
 with _bolt_runtime.scope() as _bolt_var5:
-    with _bolt_runtime.push_nesting('execute:subcommand'):
-        with _bolt_runtime.push_nesting('execute:if:score:target:targetObjective:matches:range:subcommand', *_bolt_refs[4:7]):
-            with _bolt_runtime.push_nesting('execute:commands'):
-                with _bolt_runtime.scope() as _bolt_var4:
-                    _bolt_var0 = 'thing'
-                    _bolt_var1 = 'bar'
-                    _bolt_var0 = _bolt_var0 == _bolt_var1
-                    with _bolt_helper_branch(_bolt_var0) as _bolt_condition:
-                        if _bolt_condition:
-                            _bolt_runtime.commands.extend(_bolt_refs[0].commands)
-                    _bolt_var2 = 'thing'
-                    _bolt_var3 = 'foo'
-                    _bolt_var2 = _bolt_var2 == _bolt_var3
-                    with _bolt_helper_branch(_bolt_var2) as _bolt_condition:
-                        if _bolt_condition:
-                            _bolt_runtime.commands.extend(_bolt_refs[1].commands)
+    with _bolt_runtime.push_nesting('execute:subcommand'), _bolt_runtime.push_nesting('execute:if:score:target:targetObjective:matches:range:subcommand', *_bolt_refs[4:7]), _bolt_runtime.push_nesting('execute:commands'), _bolt_runtime.scope() as _bolt_var4:
+        _bolt_var0 = 'thing'
+        _bolt_var1 = 'bar'
+        _bolt_var0 = _bolt_var0 == _bolt_var1
+        with _bolt_helper_branch(_bolt_var0) as _bolt_condition:
+            if _bolt_condition:
+                _bolt_runtime.commands.extend(_bolt_refs[0].commands)
+        _bolt_var2 = 'thing'
+        _bolt_var3 = 'foo'
+        _bolt_var2 = _bolt_var2 == _bolt_var3
+        with _bolt_helper_branch(_bolt_var2) as _bolt_condition:
+            if _bolt_condition:
+                _bolt_runtime.commands.extend(_bolt_refs[1].commands)
     _bolt_runtime.commands.append(_bolt_helper_replace(_bolt_refs[8], arguments=_bolt_helper_children([_bolt_helper_replace(_bolt_refs[7], arguments=_bolt_helper_children([*_bolt_refs[4:7], _bolt_helper_replace(_bolt_refs[3], arguments=_bolt_helper_children([_bolt_helper_replace(_bolt_refs[2], commands=_bolt_helper_children(_bolt_var4))]))]))])))
 _bolt_var6 = _bolt_helper_replace(_bolt_refs[9], commands=_bolt_helper_children(_bolt_var5))
 ---