Process review feedback

Co-authored-by: Copilot <copilot@github.com>

Author: Schamper

dissect/cstruct/cstruct.py

@@ -54,7 +54,6 @@ def __init__(self, load: str = "", *, endian: str = "<", pointer: str | None = N
         self.endian = endian
 
         self.consts = {}
-        self.lookups = {}
         self.includes = []
         # fmt: off
         self.typedefs = {

dissect/cstruct/lexer.py

@@ -27,6 +27,7 @@ class TokenType(enum.Enum):
     SEMICOLON = ";"
     COMMA = ","
     COLON = ":"
+    QUESTION = "?"
     STAR = "*"
     EQUALS = "="
 
@@ -63,7 +64,6 @@ class TokenType(enum.Enum):
     PP_FLAGS = "PP_FLAGS"
 
     # Special
-    LOOKUP = "LOOKUP"
     EOF = "EOF"
 
 
@@ -112,6 +112,7 @@ def __repr__(self) -> str:
     ";": TokenType.SEMICOLON,
     ",": TokenType.COMMA,
     ":": TokenType.COLON,
+    "?": TokenType.QUESTION,
     "*": TokenType.STAR,
     "=": TokenType.EQUALS,
     "+": TokenType.PLUS,
@@ -203,7 +204,7 @@ def _take(self, num: int = 1) -> str:
 
         return result
 
-    def _expect(self, *chars: str) -> None:
+    def _expect(self, *chars: str) -> str:
         """Consume the expected characters or raise an error."""
         if self._current() not in chars:
             actual = "end of input" if self.eof else repr(self._current())
@@ -287,11 +288,11 @@ def _read_identifier(self) -> str:
         return ""
 
     def _read_number(self) -> str:
-        """Read a numeric literal, supporting decimal, hex (0x), octal (0), binary (0b), and C-style suffixes."""
+        """Read a numeric literal, supporting decimal, hex (0x), octal (0, 0o), binary (0b), and C-style suffixes."""
         start = self._pos
         is_float = False
 
-        if self._current() == "0" and self._peek() in ("x", "X", "b", "B"):
+        if self._current() == "0" and self._peek() in ("x", "X", "b", "B", "o", "O"):
             self._expect("0")  # Consume leading 0
             suffix = self._take().lower()
 
@@ -301,6 +302,9 @@ def _read_number(self) -> str:
             if suffix == "b" and not self._read_while("01"):
                 raise self._error("invalid binary literal")
 
+            if suffix == "o" and not self._read_while("01234567"):
+                raise self._error("invalid octal literal")
+
         else:
             # Consume decimal/octal digits
             self._read_while("0123456789")
@@ -318,7 +322,9 @@ def _read_number(self) -> str:
             self._read_while("uUlL")
 
             # Convert octal: leading 0 without 0x/0b → insert 'o'
-            if len(raw) > 1 and raw[0] == "0" and raw[1].lower() not in ("x", "b"):
+            if len(raw) > 1 and raw[0] == "0" and raw[1].lower() not in ("x", "b", "o"):
+                if raw[1] not in "01234567":
+                    raise self._error("invalid octal literal")
                 raw = raw[0] + "o" + raw[1:]
 
         return raw
@@ -396,30 +402,6 @@ def _read_preprocessor(self) -> None:
 
             self._emit(TokenType.STRING, value, line)
 
-    def _read_lookup(self) -> None:
-        """Read a lookup definition: ``$name = { dict }``."""
-        line = self._line
-        col = self._column
-        start = self._pos
-
-        self._expect("$")  # Consume `$`
-
-        # Read until end of the {...} block
-        brace_depth = 0
-        while not self.eof:
-            ch = self._current()
-            if ch == "{":
-                brace_depth += 1
-            elif ch == "}":
-                brace_depth -= 1
-                if brace_depth == 0:
-                    self._expect("}")  # Consume final `}`
-                    break
-            self._take()
-
-        value = self._get(start, self._pos)
-        self._emit(TokenType.LOOKUP, value.strip(), line, col)
-
     def tokenize(self) -> list[Token]:
         """Tokenize the input data and return a list of tokens."""
         while not self.eof:
@@ -490,10 +472,6 @@ def tokenize(self) -> list[Token]:
             elif ch in _SINGLE_CHARS:
                 self._emit(_SINGLE_CHARS[ch], self._take(), line, col)
 
-            elif ch == "$":
-                # Custom lookup definition
-                self._read_lookup()
-
             else:
                 raise self._error(f"unexpected character {ch!r}", line=line)
 

dissect/cstruct/parser.py

@@ -137,8 +137,6 @@ def _parse(self) -> None:
         while (token := self._current()).type != TokenType.EOF:
             if token.type == TokenType.PP_FLAGS:
                 self._parse_config_flags()
-            elif token.type == TokenType.LOOKUP:
-                self._parse_lookup()
             elif token.type == TokenType.TYPEDEF:
                 self._parse_typedef()
             elif token.type in (TokenType.STRUCT, TokenType.UNION):
@@ -173,7 +171,7 @@ def _parse_define(self) -> None:
         while (token := self._current()).type != TokenType.EOF and token.line == name_token.line:
             parts.append(self._take().value)
 
-        value = " ".join(parts).strip()
+        value = "".join(parts).strip()
         try:
             # Lazy mode, try to evaluate as a Python literal first (for simple constants)
             value = ast.literal_eval(value)
@@ -508,17 +506,3 @@ def _parse_type_spec(self) -> type[BaseType]:
                 break
 
         return self.cs.resolve(" ".join(parts))
-
-    # Custom lookup definitions
-
-    def _parse_lookup(self) -> None:
-        """Parse a lookup definition."""
-        value = self._take().value
-
-        # Parse $name = { dict }
-        # Find the name and dict parts
-        dollar_rest = value.lstrip("$")
-        name, _, lookup = dollar_rest.partition("=")
-
-        d = ast.literal_eval(lookup.strip())
-        self.cs.lookups[name.strip()] = {self.cs.consts[k]: v for k, v in d.items()}

tests/test_basic.py

@@ -150,16 +150,6 @@ def test_typedef(cs: cstruct) -> None:
     assert cs.resolve("test") == cs.uint32
 
 
-def test_lookups(cs: cstruct, compiled: bool) -> None:
-    cdef = """
-    #define test_1 1
-    #define test_2 2
-    $a = {'test_1': 3, 'test_2': 4}
-    """
-    cs.load(cdef, compiled=compiled)
-    assert cs.lookups["a"] == {1: 3, 2: 4}
-
-
 def test_config_flag_nocompile(cs: cstruct, compiled: bool) -> None:
     cdef = """
     struct compiled_global

tests/test_expression.py

@@ -112,9 +112,9 @@ def test_expression(cs_with_consts: cstruct, expression: str, answer: int) -> No
         ),
         pytest.param(
             "$",
-            ExpressionParserError,
-            "Unmatched token: '\\$'",
-            id="invalid-token",
+            LexerError,
+            "unexpected character '\\$'",
+            id="unexpected-character",
         ),
         pytest.param(
             "-",

tests/test_lexer.py

@@ -30,6 +30,7 @@
         ("0b1010", [TokenType.NUMBER], ["0b1010"]),
         ("0B1100", [TokenType.NUMBER], ["0B1100"]),
         ("0755", [TokenType.NUMBER], ["0o755"]),
+        ("0o777", [TokenType.NUMBER], ["0o777"]),
         ("3.14", [TokenType.NUMBER], ["3.14"]),
         ("1.", [TokenType.NUMBER], ["1."]),
         ("{", [TokenType.LBRACE], ["{"]),
@@ -103,9 +104,6 @@
         ("_my_var", [TokenType.IDENTIFIER], ["_my_var"]),
         ("UINT32", [TokenType.IDENTIFIER], ["UINT32"]),
         ("uint32_t", [TokenType.IDENTIFIER], ["uint32_t"]),
-        # Lookup
-        ("$my_lookup = {1: 'a', 2: 'b'}", [TokenType.LOOKUP], ["$my_lookup = {1: 'a', 2: 'b'}"]),
-        ("$tbl = {\n    1: 'x',\n    2: 'y'\n}", [TokenType.LOOKUP], ["$tbl = {\n    1: 'x',\n    2: 'y'\n}"]),
         # Combination
         (
             "uint32_t bit0:1;",
@@ -148,6 +146,7 @@ def test_token(src: str, types: list[TokenType], values: list[str]) -> None:
     [
         ("0b", "invalid binary literal"),
         ("0x", "invalid hexadecimal literal"),
+        ("0888", "invalid octal literal"),
         ('"unterminated', "unexpected end of input"),
         ("'unterminated", "unexpected end of input"),
         ("#foobar", "unknown preprocessor directive"),

tests/test_parser.py

@@ -199,6 +199,26 @@ def test_typedef_pointer(cs: cstruct) -> None:
     assert cs.PIMAGE_DATA_DIRECTORY.type == cs._IMAGE_DATA_DIRECTORY
 
 
+def test_define(cs: cstruct) -> None:
+    cdef = """
+    #define MY_CONST 42
+    #define MY_EXPR (1 + 2 * 3)
+    #define MY_STR "hello"
+    #define MY_BYTES b"world"
+    #define MY_FUNC(x) ( x == 0 )
+    #define MY_TERNARY(x) ( x ? 1 : 0 )
+    """
+    cs.load(cdef)
+
+    assert cs.consts["MY_CONST"] == 42
+    assert cs.consts["MY_EXPR"] == 7
+    assert cs.consts["MY_STR"] == "hello"
+    assert cs.consts["MY_BYTES"] == b"world"
+    # We don't evaluate function-like macros yet, so they should be stored as their raw string representation
+    assert cs.consts["MY_FUNC"] == "(x)(x==0)"
+    assert cs.consts["MY_TERNARY"] == "(x)(x?1:0)"
+
+
 def test_undef(cs: cstruct) -> None:
     cdef = """
     #define MY_CONST 42