Small fixes in new parser

dissect/cstruct/lexer.py

@@ -214,7 +214,11 @@ def _expect(self, *chars: str) -> str:
         return self._take()
 
     def _error(self, msg: str, *, line: int | None = None) -> LexerError:
-        return LexerError(f"line {line if line is not None else self._line}: {msg}")
+        line = line if line is not None else self._line
+        context = self.data.splitlines()[line - 2 : line + 2]
+        no_len = len(str(line)) + 2
+        content = "\n".join(f"{i + line - 1:>{no_len}}: {l}" for i, l in enumerate(context))
+        return LexerError(f"line {line}: {msg}\n{line}: {content}")
 
     def _emit(self, type: TokenType, value: str, line: int, column: int = 0) -> None:
         """Emit a token with the given type and value at the specified line and column."""

dissect/cstruct/parser.py

@@ -177,8 +177,9 @@ def _parse_define(self) -> None:
                 if value[-1] != quote:
                     raise self._error("unterminated bytes literal", token=token)
 
-                # Remove the leading b and surrounding quotes
-                value = ast.literal_eval(f"b{value[2:-1]!r}")
+                # Remove the leading b and surrounding quotes and flatten escape sequences
+                value = value[2:-1].encode().decode("unicode_escape")
+                value = ast.literal_eval(f"b{value!r}")
             else:
                 try:
                     # Lazy mode, try to evaluate as a Python literal first (for simple constants)
@@ -385,7 +386,15 @@ def _parse_enum_or_flag(self) -> type[Enum | Flag]:
                 continue
             self._assert_not_eof()
 
-            member_name = self._expect(TokenType.IDENTIFIER).value
+            # For historical reasons, we allow enum/flag member names to start with a digit
+            # E.g. `32BIT`
+            member_name = ""
+            if token := self._match(TokenType.NUMBER):
+                member_name += token.value
+                if token := self._match(TokenType.IDENTIFIER):
+                    member_name += token.value
+            else:
+                member_name = self._expect(TokenType.IDENTIFIER).value
 
             if self._match(TokenType.EQUALS):
                 expression = self._collect_until(TokenType.COMMA, TokenType.RBRACE)

tests/test_lexer.py

@@ -182,6 +182,18 @@ def test_error(src: str, match: str) -> None:
         tokenize(src)
 
 
+def test_error_context() -> None:
+    """Test that a LexerError includes the surrounding lines as context."""
+    src = "struct test {\n    uint32 a;\n    @invalid;\n    uint32 b;\n};"
+    with pytest.raises(LexerError, match="line 3: unexpected character '@'") as exc_info:
+        tokenize(src)
+
+    message = str(exc_info.value)
+    assert "  2:     uint32 a;" in message
+    assert "  3:     @invalid;" in message
+    assert "  4:     uint32 b;" in message
+
+
 def test_line_and_column_tracking() -> None:
     """Test that the lexer correctly tracks line and column numbers."""
     src = "a\n  b\nc"

tests/test_parser.py

@@ -214,6 +214,27 @@ def test_typedef_enum(cs: cstruct) -> None:
     assert cs.test_enum.VAL3 == 4
 
 
+def test_enum_flag_digit_member_name(cs: cstruct) -> None:
+    # For historical reasons, we allow enum/flag member names to start with a digit, e.g. `32BIT`
+    cdef = """
+    enum test_enum : uint8 {
+        32BIT = 1,
+        64BIT = 2
+    };
+
+    flag test_flag : uint8 {
+        32BIT = 1,
+        64BIT = 2
+    };
+    """
+    cs.load(cdef)
+
+    assert cs.test_enum["32BIT"] == 1
+    assert cs.test_enum["64BIT"] == 2
+    assert cs.test_flag["32BIT"] == 1
+    assert cs.test_flag["64BIT"] == 2
+
+
 def test_define(cs: cstruct) -> None:
     cdef = """
     #define CONST 42
@@ -230,6 +251,7 @@ def test_define(cs: cstruct) -> None:
                         3)
     #define QUOTES "\'\"a'b\""
     #define ESCAPE "\\'\\"a'b\\"\\n"
+    #define BYTES_ESCAPE b"`\\n"
     #define FUNC(x) ( x == 0 )
     #define TERNARY(x) ( x ? 1 : 0 )
     """
@@ -247,6 +269,7 @@ def test_define(cs: cstruct) -> None:
     assert cs.consts["MULTILINE"] == 6
     assert cs.consts["QUOTES"] == "'\"a'b\""
     assert cs.consts["ESCAPE"] == "'\"a'b\"\n"
+    assert cs.consts["BYTES_ESCAPE"] == b"`\n"
     # We don't evaluate function-like macros yet, so they should be stored as their raw string representation
     assert cs.consts["FUNC"] == "(x) ( x == 0 )"
     assert cs.consts["TERNARY"] == "(x) ( x ? 1 : 0 )"