Add parser test and error for free code outside entry

jac/jaclang/compiler/__init__.py

@@ -63,14 +63,18 @@ def gen_all_parsers() -> None:
 
 
 try:
-    from jaclang.compiler.larkparse import jac_parser as jac_lark
+    from jaclang.compiler.larkparse import (  # type: ignore[attr-defined]
+        jac_parser as jac_lark,
+    )
 
     jac_lark.logger.setLevel(logging.DEBUG)
 except (ModuleNotFoundError, ImportError):
     print("Parser not present, generating for developer setup...", file=sys.stderr)
     try:
         gen_all_parsers()
-        from jaclang.compiler.larkparse import jac_parser as jac_lark
+        from jaclang.compiler.larkparse import (  # type: ignore[attr-defined]
+            jac_parser as jac_lark,
+        )
 
         jac_lark.logger.setLevel(logging.DEBUG)
     except Exception as e:

jac/jaclang/compiler/parser.py

@@ -184,6 +184,17 @@ def feed_current_token(iparser: jl.InteractiveParser, tok: jl.Token) -> bool:
                 iparser.feed_token(jl.Token(Tok.NAME.name, "recover_name_token"))
                 return feed_current_token(iparser, e.token)
 
+            if Tok.KW_WITH.name in iparser.accepts():
+                # Copy and feed the `with` token cause we don't want to actually
+                # feed the `with` token in the main parser state.
+                ipcopy = iparser.copy(deepcopy_values=True)
+                ipcopy.feed_token(jl.Token(Tok.KW_WITH.name, "with"))
+                if Tok.KW_ENTRY.name in ipcopy.accepts():
+                    self.log_error(
+                        "Arbitary statements must be declared inside an entry block",
+                        self.error_to_token(e),
+                    )
+
             # We're calling try_feed_missing_token twice here because the first missing
             # will be reported as such and we don't for the consequent missing token.
             if tk := try_feed_missing_token(iparser):

jac/jaclang/compiler/tsparser.py

@@ -11,7 +11,7 @@
 
 import jaclang.compiler.unitree as uni
 from jaclang.compiler.constant import TsTokens as Tok
-from jaclang.compiler.larkparse import ts_parser as ts_lark
+from jaclang.compiler.larkparse import ts_parser as ts_lark  # type: ignore[attr-defined]
 from jaclang.compiler.passes.main import BaseTransform, Transform
 
 if TYPE_CHECKING:

jac/jaclang/langserve/tests/server_test/test_lang_serve.py

@@ -63,7 +63,10 @@ def test_open_with_syntax_error():
 
     try:
         helper.open_document()
-        helper.assert_has_diagnostics(count=2, message_contains="Unexpected token")
+        helper.assert_has_diagnostics(
+            count=3,
+            message_contains="Arbitary statements must be declared inside an entry block",
+        )
 
         diagnostics = helper.get_diagnostics()
         assert str(diagnostics[0].range) == "59:0-59:5"
@@ -94,7 +97,7 @@ def test_did_open_and_simple_syntax_error():
             test_file._get_template_path(CIRCLE_TEMPLATE), "error"
         )
         helper.change_document(broken_code)
-        helper.assert_has_diagnostics(count=2)
+        helper.assert_has_diagnostics(count=3)
         helper.assert_semantic_tokens_count(EXPECTED_CIRCLE_TOKEN_COUNT_ERROR)
     finally:
         ls.shutdown()
@@ -124,7 +127,7 @@ def test_did_save():
         )
         helper.save_document(broken_code)
         helper.assert_semantic_tokens_count(EXPECTED_CIRCLE_TOKEN_COUNT_ERROR)
-        helper.assert_has_diagnostics(count=2, message_contains="Unexpected token")
+        helper.assert_has_diagnostics(count=3, message_contains="Unexpected token")
     finally:
         ls.shutdown()
         test_file.cleanup()
@@ -153,7 +156,7 @@ def test_did_change():
         helper.change_document("\nerror" + test_file.code)
         helper.assert_semantic_tokens_count(EXPECTED_CIRCLE_TOKEN_COUNT)
         helper.assert_has_diagnostics(
-            count=2, message_contains="Unexpected token 'error'"
+            count=3, message_contains="Unexpected token 'error'"
         )
     finally:
         ls.shutdown()
@@ -203,7 +206,7 @@ def test_multifile_workspace():
 
         # Verify initial state
         helper1.assert_no_diagnostics()
-        helper2.assert_has_diagnostics(count=1, message_contains="Unexpected token")
+        helper2.assert_has_diagnostics(count=2, message_contains="Unexpected token")
 
         # Check semantic tokens before change
         helper1.assert_semantic_tokens_count(EXPECTED_GLOB_TOKEN_COUNT)

jac/jaclang/tests/test_language.py

@@ -6,7 +6,7 @@
 import sys
 import tempfile
 from collections.abc import Callable, Generator
-from contextlib import AbstractContextManager
+from contextlib import AbstractContextManager, redirect_stderr
 from pathlib import Path
 from unittest.mock import patch
 
@@ -231,6 +231,51 @@ def test_assignment_list_no_infinite_loop():
     assert len(jac_prog.errors_had) > 0  # Check errors on program
 
 
+def test_free_code_outside_entry_reports_error(
+    fixture_path: Callable[[str], str],
+    capture_stdout: Callable[[], AbstractContextManager[io.StringIO]],
+) -> None:
+    """Free python code outside an entry should report an error."""
+
+    # Case 1: Free python code not wrapped in `with entry { ... }` is invalid
+    prog = JacProgram()
+    code = 'print("hello world");'
+
+    # Compile and capture stderr via pytest capsys
+    stderr_buffer = io.StringIO()
+    with redirect_stderr(stderr_buffer) as captured_output:
+        prog.compile(use_str=code, file_path="test.jac")
+    captured = captured_output.getvalue()
+
+    # The parser change logs this specific message when encountering free code
+    assert "Arbitary statements must be declared inside an entry block" in captured
+
+    # Case 2: This is unexpected token as well but not the same error
+    prog = JacProgram()
+    code = """
+    with entry {
+        *;
+    }
+    """
+    stderr_buffer = io.StringIO()
+    with redirect_stderr(stderr_buffer) as captured_output:
+        prog.compile(use_str=code, file_path="test.jac")
+    captured = captured_output.getvalue()
+    assert "Arbitary statements must be declared inside an entry block" not in captured
+
+    # Case 3: `with` context manager block
+    prog = JacProgram()
+    code = """
+    with entry {
+        with open('some/file.txt', 'r') as f {
+            print(f.read());
+        }
+    """
+    with capture_stdout():
+        prog.compile(use_str=code, file_path="test.jac")
+    assert len(prog.errors_had) == 0
+
+
 def test_need_import(
     fixture_path: Callable[[str], str],
     capture_stdout: Callable[[], AbstractContextManager[io.StringIO]],