fix: only allow imports and macros directly at scope level

bolt/parse.py

@@ -44,6 +44,7 @@
     "parse_name_list",
     "DeferredRootBacktracker",
     "FunctionConstraint",
+    "RootScopeHandler",
     "BinaryParser",
     "UnaryParser",
     "UnpackParser",
@@ -534,6 +535,7 @@ def create_bolt_root_parser(parser: Parser, macro_handler: "MacroHandler"):
     parser = DeferredRootBacktracker(parser, macro_handler=macro_handler)
     parser = DecoratorResolver(parser)
     parser = ProcMacroExpansion(parser)
+    parser = RootScopeHandler(parser)
     return parser
 
 
@@ -992,6 +994,7 @@ def parse_deferred_root(stream: TokenStream) -> AstDeferredRoot:
     stream_copy.data["deferred_locals"] = set()
     stream_copy.data["branch_scope"] = set()
     stream_copy.data["class_scope"] = None
+    del stream_copy.data["root_scope"]
 
     deferred_locals.clear()
 
@@ -1139,6 +1142,7 @@ def __call__(self, stream: TokenStream) -> Any:
             return node
 
         if node.identifier == "macro:name:subcommand":
+            self.check_root_scope(stream, node)
             node = self.create_macro(node, stream.data.get("resource_location"))
             if not isinstance(node, AstProcMacro):
                 declaration = replace(node, arguments=AstChildren(node.arguments[:-1]))
@@ -1153,6 +1157,11 @@ def __call__(self, stream: TokenStream) -> Any:
 
         return node
 
+    def check_root_scope(self, stream: TokenStream, node: AstCommand):
+        if not stream.data.get("root_scope"):
+            message = "Macro definition can only appear directly at scope level."
+            raise set_location(InvalidSyntax(message), node, node.arguments[0])
+
     def create_macro(
         self,
         command: AstCommand,
@@ -1474,11 +1483,18 @@ class ImportStatementHandler:
     def __call__(self, stream: TokenStream) -> Any:
         if isinstance(node := self.parser(stream), AstCommand):
             if node.identifier in ["import:module", "import:module:as:alias"]:
+                self.check_root_scope(stream, node)
                 return self.handle_import(stream, node)
             elif node.identifier == "from:module:import:subcommand":
+                self.check_root_scope(stream, node)
                 return self.handle_from_import(stream, node)
         return node
 
+    def check_root_scope(self, stream: TokenStream, node: AstCommand):
+        if not stream.data.get("root_scope"):
+            message = "Import statement can only appear directly at scope level."
+            raise set_location(InvalidSyntax(message), node)
+
     def handle_import(self, stream: TokenStream, node: AstCommand) -> AstCommand:
         identifiers = get_stream_identifiers(stream)
         identifiers_storage = get_stream_identifiers_storage(stream)
@@ -1717,6 +1733,17 @@ def __call__(self, stream: TokenStream) -> AstRoot:
         return node
 
 
+@dataclass
+class RootScopeHandler:
+    """Handle root scope."""
+
+    parser: Parser
+
+    def __call__(self, stream: TokenStream) -> Any:
+        with stream.provide(root_scope="root_scope" not in stream.data):
+            return self.parser(stream)
+
+
 @dataclass
 class BinaryParser:
     """Parser for binary expressions."""

tests/resources/bolt_examples.mcfunction

@@ -990,3 +990,30 @@ class B:
     def a(self):
         say a
         return A()
+###
+if 0:
+    import math
+###
+def f():
+    if 0:
+        import this
+###
+for i in range(10):
+    if i % 2 == 0:
+        from math import sin
+###
+if "beans":
+    macro beans:
+        pass
+###
+for bean in "beans":
+    if bean:
+        macro jelly bean(stream):
+            pass
+###
+class A:
+    macro foo:
+        pass
+###
+class B:
+    import math

tests/snapshots/bolt__parse_266__0.txt

@@ -0,0 +1,7 @@
+if 0:
+#>ERROR  Import statement can only appear directly at scope level.
+# line 2, column 5
+#      1 |  if 0:
+#      2 |      import math
+#        :      ^^^^^^^^^^^
+    import math

tests/snapshots/bolt__parse_267__0.txt

@@ -0,0 +1,8 @@
+def f():
+    if 0:
+#>ERROR  Import statement can only appear directly at scope level.
+# line 3, column 9
+#      2 |      if 0:
+#      3 |          import this
+#        :          ^^^^^^^^^^^
+        import this

tests/snapshots/bolt__parse_268__0.txt

@@ -0,0 +1,8 @@
+for i in range(10):
+    if i % 2 == 0:
+#>ERROR  Import statement can only appear directly at scope level.
+# line 3, column 9
+#      2 |      if i % 2 == 0:
+#      3 |          from math import sin
+#        :          ^^^^^^^^^^^^^^^^^^^^
+        from math import sin

tests/snapshots/bolt__parse_269__0.txt

@@ -0,0 +1,9 @@
+if "beans":
+#>ERROR  Macro definition can only appear directly at scope level.
+# line 2, column 5
+#      1 |  if "beans":
+#      2 |      macro beans:
+#        :      ^^^^^^^^^^^
+#      3 |          pass
+    macro beans:
+        pass

tests/snapshots/bolt__parse_270__0.txt

@@ -0,0 +1,10 @@
+for bean in "beans":
+    if bean:
+#>ERROR  Macro definition can only appear directly at scope level.
+# line 3, column 9
+#      2 |      if bean:
+#      3 |          macro jelly bean(stream):
+#        :          ^^^^^^^^^^^
+#      4 |              pass
+        macro jelly bean(stream):
+            pass

tests/snapshots/bolt__parse_271__0.txt

@@ -0,0 +1,9 @@
+class A:
+#>ERROR  Macro definition can only appear directly at scope level.
+# line 2, column 5
+#      1 |  class A:
+#      2 |      macro foo:
+#        :      ^^^^^^^^^
+#      3 |          pass
+    macro foo:
+        pass

tests/snapshots/bolt__parse_272__0.txt

@@ -0,0 +1,7 @@
+class B:
+#>ERROR  Import statement can only appear directly at scope level.
+# line 2, column 5
+#      1 |  class B:
+#      2 |      import math
+#        :      ^^^^^^^^^^^
+    import math