dissect.cstruct 5.0.dev10__py3-none-any.whl → 5.0.dev12__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -5,10 +5,14 @@ from __future__ import annotations
5
5
  import io
6
6
  import logging
7
7
  from enum import Enum
8
- from textwrap import dedent, indent
9
- from typing import TYPE_CHECKING
10
-
11
- from dissect.cstruct.bitbuffer import BitBuffer
8
+ from itertools import groupby, repeat
9
+ from keyword import iskeyword
10
+ from operator import itemgetter
11
+ from textwrap import indent
12
+ from typing import TYPE_CHECKING, Any
13
+
14
+ from dissect.cstruct.exception import ExpressionParserError
15
+ from dissect.cstruct.expression import SOURCE_GLOBALS, Expression
12
16
  from dissect.cstruct.types import (
13
17
  Array,
14
18
  BaseArray,
@@ -26,13 +30,14 @@ from dissect.cstruct.types import (
26
30
  WcharArray,
27
31
  )
28
32
  from dissect.cstruct.types.enum import EnumMetaType
29
- from dissect.cstruct.util import _struct
33
+ from dissect.cstruct.types.structure import UnionMetaType, UnionProxy
34
+ from dissect.cstruct.util import ENDIANNESS_TO_BYTEORDER_MAP, _struct
30
35
 
31
36
  if TYPE_CHECKING:
32
- from collections.abc import Iterator
37
+ from collections.abc import Callable, Iterable, Iterator
33
38
  from types import MethodType
39
+ from typing import BinaryIO
34
40
 
35
- from dissect.cstruct.cstruct import cstruct
36
41
  from dissect.cstruct.types import BaseType
37
42
  from dissect.cstruct.types.structure import Field
38
43
 
@@ -57,95 +62,295 @@ log = logging.getLogger(__name__)
57
62
  python_compile = compile
58
63
 
59
64
 
60
- def compile(structure: type[Structure]) -> type[Structure]:
61
- return Compiler(structure.__cs__).compile(structure)
65
+ def compile_structure(cls: type[Structure]) -> type[Structure]:
66
+ """Compile the ``_read`` method of a structure or union type, if the compiler supports it.
62
67
 
68
+ The compilation happens on first use.
69
+ """
70
+ try:
71
+ _check_supported(cls, cls.__fields__)
72
+ cls._read = _LazyCompileRead(cls)
73
+ cls.__compiled__ = True
74
+ except Exception as e:
75
+ # Revert to the interpreted reader
76
+ log.debug("Failed to compile %s", cls, exc_info=e)
77
+ cls._read = classmethod(type(cls)._read)
78
+ cls.__compiled__ = False
63
79
 
64
- class Compiler:
65
- def __init__(self, cs: cstruct):
66
- self.cs = cs
80
+ return cls
67
81
 
68
- def compile(self, structure: type[Structure]) -> type[Structure]:
69
- if issubclass(structure, Union):
70
- return structure
71
82
 
72
- try:
73
- structure._read = self.compile_read(structure.__fields__, structure.__name__, structure.__align__)
74
- structure.__compiled__ = True
75
- except Exception as e:
76
- # Silently ignore, we didn't compile unfortunately
77
- log.debug("Failed to compile %s", structure, exc_info=e)
83
+ class _LazyCompileRead:
84
+ """Descriptor that compiles the ``_read`` method of a structure on first access and replaces itself with it."""
78
85
 
79
- return structure
86
+ def __init__(self, cls: type[Structure]):
87
+ self.cls = cls
88
+ self.fields = list(cls.__fields__)
89
+ self.method = None
80
90
 
81
- def compile_read(self, fields: list[Field], name: str | None = None, align: bool = False) -> MethodType:
82
- return _ReadSourceGenerator(self.cs, fields, name, align).generate()
91
+ # Snapshot of the nested structure types, so that we only inline them if they're unchanged on first use
92
+ self.snapshot = _snapshot_structures(self.fields, {})
83
93
 
94
+ def __get__(self, obj: Any, owner: type | None = None) -> MethodType:
95
+ return self.compile().__get__(obj, owner)
84
96
 
85
- class _ReadSourceGenerator:
86
- def __init__(self, cs: cstruct, fields: list[Field], name: str | None = None, align: bool = False):
87
- self.cs = cs
88
- self.fields = fields
89
- self.name = name
90
- self.align = align
97
+ def compile(self) -> classmethod:
98
+ if self.method is None:
99
+ cls = self.cls
100
+
101
+ # Since we lazily compile, some of the nested structures may have changed
102
+ # Only inline if they're the same as when we last built the structure
103
+ inline = all(type_.__dict__.get("__init__") is init for type_, init in self.snapshot.items())
104
+
105
+ try:
106
+ self.method = _ReadSourceGenerator(cls, self.fields, inline).generate()
107
+ except Exception as e:
108
+ # Revert to the slower loop based method
109
+ log.debug("Failed to compile %s", cls, exc_info=e)
110
+ self.method = classmethod(type(cls)._read)
111
+ cls.__compiled__ = False
112
+
113
+ if cls.__dict__.get("_read") is self:
114
+ cls._read = self.method
91
115
 
92
- self.field_map: dict[str, Field] = {}
93
- self._token_id = 0
116
+ return self.method
94
117
 
95
- def _map_field(self, field: Field) -> str:
96
- token = f"_{self._token_id}"
97
- self.field_map[token] = field
98
- self._token_id += 1
99
- return token
118
+
119
+ def _snapshot_structures(fields: list[Field], result: dict[type, Any]) -> dict[type, Any]:
120
+ for field in fields:
121
+ type_ = field.type
122
+ if issubclass(type_, Structure) and type_ not in result:
123
+ result[type_] = type_.__dict__.get("__init__")
124
+ _snapshot_structures(type_.__fields__, result)
125
+ return result
126
+
127
+
128
+ def _check_supported(cls: type[Structure], fields: list[Field]) -> None:
129
+ """Check if the compiler supports the given fields."""
130
+ if issubclass(cls, Union):
131
+ _check_union(cls, fields)
132
+ return
133
+
134
+ for field in fields:
135
+ field_type = field.type
136
+
137
+ if isinstance(field_type, EnumMetaType):
138
+ field_type = field_type.__type__
139
+
140
+ if not issubclass(field_type, SUPPORTED_TYPES):
141
+ raise TypeError(f"Unsupported type for compiler: {field_type}")
142
+
143
+ if field.bits and field_type.__dynamic__:
144
+ raise TypeError(f"Unsupported type for bit field: {field_type}")
145
+
146
+
147
+ class _Fallback:
148
+ """Lazily compiled ``_read`` function without inlined structures.
149
+
150
+ Used when one of the inlined structure types was modified after compilation.
151
+ """
152
+
153
+ def __init__(self, cls: type[Structure], fields: list[Field]):
154
+ self.cls = cls
155
+ self.fields = list(fields)
156
+ self.func = None
157
+
158
+ def __call__(self, cls: type[Structure], stream: BinaryIO, context: dict[str, Any] | None, endian: str) -> Any:
159
+ if self.func is None:
160
+ self.func = _ReadSourceGenerator(self.cls, self.fields, inline=False).generate().__func__
161
+ return self.func(cls, stream, context=context, endian=endian)
162
+
163
+
164
+ _ENDIANNESS = ("<", ">", "!", "@", "=")
165
+
166
+ # Types of which the value can be used as an integer directly
167
+ _INTEGER_TYPES = (Int, Packed, Enum, Flag)
168
+
169
+ # Globals that are available to all generated read functions
170
+ _GLOBALS = {
171
+ "_type_call": type.__call__,
172
+ "_object_new": object.__new__,
173
+ "_setattr": setattr,
174
+ "_repeat": repeat,
175
+ "_int_from_bytes": int.from_bytes,
176
+ "_byteorder": ENDIANNESS_TO_BYTEORDER_MAP,
177
+ "_object_setattr": object.__setattr__,
178
+ "_getattr": getattr,
179
+ "_UnionProxy": UnionProxy,
180
+ "_union_read": UnionMetaType._read,
181
+ **SOURCE_GLOBALS,
182
+ }
183
+
184
+
185
+ class _ReadSourceGenerator:
186
+ """Generates the source code of an optimized ``_read`` method for a structure or union."""
187
+
188
+ def __init__(self, cls: type[Structure], fields: list[Field], inline: bool = True):
189
+ self.cls = cls
190
+ self.fields = fields
191
+ self.align = cls.__align__
192
+ self.inline = inline and not self.align
193
+
194
+ # Witnesses of inlined structure types (symbol of the type -> symbol of its __init__ at compile time)
195
+ self._witnesses: dict[str, str] = {}
196
+ self._fallback_symbol: str | None = None
197
+
198
+ # The state of the bit buffer: the type of its current value and the remaining bits
199
+ self._bits: tuple[type[BaseType], int] | None = None
200
+ self._var_ids: dict[str, int] = {}
201
+
202
+ # The values available to the generated code
203
+ self.symbols: dict[str, Any] = {"cs": cls.__cs__}
204
+ self._symbol_names: dict[int, str] = {}
205
+
206
+ # (field name, local variable name) of all fields assigned so far
207
+ self._assigned: list[tuple[str, str]] = []
208
+ self._has_context = False
209
+ self._uses_start = False
210
+ self._uses_sizes = False
211
+
212
+ def _type_symbol(self, type_: type[BaseType]) -> str:
213
+ """Add a type as symbol, returns its name."""
214
+ return self._add_symbol(type_, "type")
215
+
216
+ def _add_symbol(self, value: Any, name: str) -> str:
217
+ """Add a value that is available to the generated code, returns its name."""
218
+ if (symbol := self._symbol_names.get(id(value))) is None:
219
+ symbol = self._symbol_names[id(value)] = self._var(name)
220
+ self.symbols[symbol] = value
221
+ return symbol
222
+
223
+ def _context(self, lines: list[str]) -> str:
224
+ """Return the name of the context dictionary, creating it from the fields read so far if needed."""
225
+ if not self._has_context:
226
+ self._has_context = True
227
+ items = ", ".join(f"{name!r}: {var}" for name, var in self._assigned)
228
+ lines.append(f"values = {{{items}}}")
229
+
230
+ return "values"
231
+
232
+ def _var(self, name: str) -> str:
233
+ """Return a unique variable name for the generated code."""
234
+ index = self._var_ids.get(name, 0)
235
+ self._var_ids[name] = index + 1
236
+ return f"_{name}{index}"
237
+
238
+ def _assign(self, field: Field, value: str) -> str:
239
+ """Assign a value to a new local variable for the given field, and the context dictionary if it exists."""
240
+ var = f"_value{len(self._assigned)}"
241
+ self._assigned.append((field._name, var))
242
+ if self._has_context:
243
+ return f"{var} = values[{field._name!r}] = {value}"
244
+ return f"{var} = {value}"
245
+
246
+ def _assign_local(self, field: Field, value: str) -> list[str]:
247
+ return [self._assign(field, value)]
248
+
249
+ def _fallback(self) -> str:
250
+ """Return the symbol of the non-inlined fallback read function."""
251
+ if self._fallback_symbol is None:
252
+ fallback = _Fallback(self.cls, self.fields)
253
+ self._fallback_symbol = self._add_symbol(fallback, "fallback")
254
+ return self._fallback_symbol
255
+
256
+ def _witness(self, type_: type[Structure]) -> str:
257
+ """Add an inlined type as symbol, with a witness to detect modifications to it after compilation."""
258
+ symbol = self._type_symbol(type_)
259
+ self._witnesses[symbol] = self._add_symbol(type_.__init__, "init")
260
+ return symbol
261
+
262
+ def _modified_checks(self) -> list[str]:
263
+ """Return the conditions that are true if any of the inlined types were modified after compilation."""
264
+ return [f"{type_}.__init__ is not {init}" for type_, init in self._witnesses.items()]
100
265
 
101
266
  def generate(self) -> MethodType:
102
267
  source = self.generate_source()
103
- symbols = {token: field.type for token, field in self.field_map.items()}
104
268
 
105
- code = python_compile(source, f"<compiled {self.name or 'anonymous'}._read>", "exec")
106
- exec(code, {"BitBuffer": BitBuffer, "_struct": _struct, **symbols}, d := {})
269
+ code = python_compile(source, f"<compiled {self.cls.__name__}._read>", "exec")
270
+ exec(code, {**_GLOBALS, **self.symbols}, d := {})
107
271
  obj = d.popitem()[1]
108
272
  obj.__source__ = source
109
273
 
110
274
  return classmethod(obj)
111
275
 
112
276
  def generate_source(self) -> str:
113
- preamble = """
114
- r = {}
115
- s = {}
116
- o = stream.tell()
117
- """
277
+ if issubclass(self.cls, Union):
278
+ return self._generate_union_source()
279
+ return self._generate_struct_source()
280
+
281
+ def _generate_struct_source(self) -> str:
282
+ """Generate the code for reading a structure."""
283
+ read_code = "\n".join(self._generate_fields_source())
284
+
285
+ preamble = []
286
+ if checks := self._modified_checks():
287
+ # Fall back to the non-inlined version if any of the inlined types was modified after compilation
288
+ preamble.append(f"if {' or '.join(checks)}: return {self._fallback()}(cls, stream, context, endian)")
289
+ if self._uses_sizes:
290
+ preamble.append("sizes = {}")
291
+ if self._uses_start:
292
+ preamble.append("start = stream.tell()")
293
+
294
+ sizes = "sizes" if self._uses_sizes else "{}"
295
+ kwargs = ", ".join(f"{name!r}: {var}" for name, var in self._assigned)
296
+ outro = []
297
+
298
+ if self._can_construct_directly():
299
+ # Construct the object directly if we're reading the exact structure class we compiled for
300
+ # This bypasses the __init__ method
301
+ outro.append(f"if cls is {self._add_symbol(self.cls, 'cls')}:")
302
+ outro.append(" obj = _object_new(cls)")
303
+ for name, var in self._assigned:
304
+ outro.append(f" {_set_attribute('obj', name, var)}")
305
+ outro.append(f" obj.__dynamic_sizes__ = {sizes}")
306
+ outro.append(" return obj")
307
+
308
+ # Using type.__call__ directly calls the __init__ method of the class
309
+ outro.append(f"obj = _type_call(cls, **{{{kwargs}}})")
310
+ outro.append(f"obj.__dynamic_sizes__ = {sizes}")
311
+ outro.append("return obj")
312
+
313
+ code = indent("\n".join([*preamble, read_code, *outro]), " ")
118
314
 
119
- if any(field.bits for field in self.fields):
120
- preamble += "bit_reader = BitBuffer(stream, endian=endian)\n"
315
+ return f"def _read(cls, stream, *, context=None, endian):\n{code}"
121
316
 
122
- read_code = "\n".join(self._generate_fields())
317
+ def _can_construct_directly(self) -> bool:
318
+ if not _uses_generated_init(self.cls):
319
+ return False
123
320
 
124
- outro = """
125
- obj = type.__call__(cls, **r)
126
- obj.__dynamic_sizes__ = s
321
+ # Void fields are not read and rely on the default value set by __init__
322
+ return not any(issubclass(_get_read_type(field.type), (Void, VoidArray)) for field in self.fields)
127
323
 
128
- return obj
129
- """
324
+ def _generate_union_source(self) -> str:
325
+ """Generate the code for reading a union."""
326
+ size = _check_union(self.cls, self.fields, self.inline)
327
+
328
+ lines = ["obj = _object_new(cls)"]
329
+ self._pack_union("obj", self.fields, 0, "buf", lines)
330
+
331
+ # Fall back to the generic method if any of the inlined types was modified after compilation
332
+ checks = [f"cls is not {self._add_symbol(self.cls, 'cls')}", *self._modified_checks()]
130
333
 
131
- code = indent(dedent(preamble).lstrip() + read_code + dedent(outro), " ")
334
+ preamble = [
335
+ f"if {' or '.join(checks)}: return _union_read(cls, stream, context=context, endian=endian)",
336
+ f"buf = stream.read({size})",
337
+ _check_eof(size),
338
+ ]
339
+ code = indent("\n".join([*preamble, *lines, "return obj"]), " ")
132
340
 
133
341
  return f"def _read(cls, stream, *, context=None, endian):\n{code}"
134
342
 
135
- def _generate_fields(self) -> Iterator[str]:
343
+ def _generate_fields_source(self) -> Iterator[str]:
344
+ """Generate the code for reading all fields."""
345
+ _check_supported(self.cls, self.fields)
346
+
347
+ # The offset of the current field in the structure, None if it's unknown due to a dynamic field
136
348
  current_offset = 0
137
349
  current_block: list[Field] = []
138
- prev_was_bits = False
139
- prev_bits_type = None
140
- bits_remaining = 0
141
- bits_rollover = False
142
350
 
143
351
  def flush() -> Iterator[str]:
144
352
  if current_block:
145
- if self.align and current_block[0].offset is None:
146
- yield f"stream.seek(-stream.tell() & ({current_block[0].alignment} - 1), {io.SEEK_CUR})"
147
-
148
- yield from self._generate_packed(current_block)
353
+ yield from self._generate_packed_source(current_block)
149
354
  current_block[:] = []
150
355
 
151
356
  def align_to_field(field: Field) -> Iterator[str]:
@@ -153,7 +358,8 @@ class _ReadSourceGenerator:
153
358
 
154
359
  if field.offset is not None and field.offset != current_offset:
155
360
  # If a field has a set offset and it's not the same as the current tracked offset, seek to it
156
- yield f"stream.seek(o + {field.offset})"
361
+ self._uses_start = True
362
+ yield f"stream.seek(start + {field.offset})"
157
363
  current_offset = field.offset
158
364
 
159
365
  if self.align and field.offset is None:
@@ -165,186 +371,447 @@ class _ReadSourceGenerator:
165
371
  if isinstance(field_type, EnumMetaType):
166
372
  field_type = field_type.__type__
167
373
 
168
- if not issubclass(field_type, SUPPORTED_TYPES):
169
- raise TypeError(f"Unsupported type for compiler: {field_type}")
374
+ if not field.bits:
375
+ self._bits = None
170
376
 
171
- if prev_was_bits and not field.bits:
172
- yield "bit_reader.reset()"
173
- prev_was_bits = False
174
- bits_remaining = 0
377
+ size = field_type.__size__
175
378
 
176
- try:
177
- size = len(field_type)
178
- is_dynamic = False
179
- except TypeError:
180
- size = None
181
- is_dynamic = True
182
-
183
- # Sub structure
184
- if issubclass(field_type, Structure):
185
- yield from flush()
186
- yield from align_to_field(field)
187
- yield from self._generate_structure(field)
379
+ # Sub structure that we can inline into the current block
380
+ inline = (
381
+ self.inline
382
+ and field.offset is not None
383
+ and field.offset == current_offset
384
+ and _inline_kind(field_type) is not None
385
+ )
188
386
 
189
- # Array of structures and multi-dimensional arrays
190
- elif issubclass(field_type, (Array, CharArray, WcharArray)) and (
191
- issubclass(field_type.__type__, Structure) or issubclass(field_type.__type__, BaseArray) or is_dynamic
387
+ # Sub structures, arrays of structures, multi-dimensional and dynamic arrays
388
+ if (issubclass(field_type, Structure) and not inline) or (
389
+ issubclass(field_type, (Array, CharArray, WcharArray))
390
+ and (issubclass(field_type.__type__, (Structure, BaseArray)) or size is None)
192
391
  ):
193
392
  yield from flush()
194
393
  yield from align_to_field(field)
195
- yield from self._generate_array(field)
394
+ yield from self._generate_field_source(field)
196
395
 
197
396
  # Bit fields
198
397
  elif field.bits:
199
- if size is None:
200
- raise TypeError(f"Unsupported type for bit field: {field_type}")
201
-
202
- if not prev_was_bits:
203
- prev_bits_type = field_type
204
- prev_was_bits = True
205
-
206
- if bits_remaining == 0 or prev_bits_type != field_type:
207
- bits_remaining = (size * 8) - field.bits
208
- bits_rollover = True
398
+ if not self._starts_bit_value(field):
399
+ # Only reading a new value into the bit buffer moves the stream
400
+ size = 0
209
401
 
210
402
  yield from flush()
211
403
  yield from align_to_field(field)
212
- yield from self._generate_bits(field)
404
+ yield from self._generate_bits_source(field)
213
405
 
214
- # Everything else - basic and composite types (and arrays of them)
406
+ # Everything else - basic and composite types (and arrays of them), and inlined structures
215
407
  else:
408
+ if not current_block:
409
+ yield from align_to_field(field)
216
410
  current_block.append(field)
217
411
 
218
- if current_offset is not None and size is not None and (not field.bits or bits_rollover):
219
- current_offset += size
220
- bits_rollover = False
412
+ if current_offset is not None:
413
+ current_offset = None if size is None else current_offset + size
221
414
 
222
415
  yield from flush()
223
416
 
224
417
  if self.align:
225
418
  yield f"stream.seek(-stream.tell() & (cls.__alignment__ - 1), {io.SEEK_CUR})"
226
419
 
227
- def _generate_structure(self, field: Field) -> Iterator[str]:
228
- template = f"""
229
- {"_s = stream.tell()" if field.type.__dynamic__ else ""}
230
- r["{field._name}"] = {self._map_field(field)}._read(stream, context=r, endian=endian)
231
- {f's["{field._name}"] = stream.tell() - _s' if field.type.__dynamic__ else ""}
232
- """
233
-
234
- yield dedent(template)
420
+ def _generate_field_source(self, field: Field) -> Iterator[str]:
421
+ """Generate the code for reading and assigning a field."""
422
+ lines = []
423
+ read = self._generate_field_read_source(field, lines)
235
424
 
236
- def _generate_array(self, field: Field) -> Iterator[str]:
237
- template = f"""
238
- {"_s = stream.tell()" if field.type.__dynamic__ else ""}
239
- r["{field._name}"] = {self._map_field(field)}._read(stream, context=r, endian=endian)
240
- {f's["{field._name}"] = stream.tell() - _s' if field.type.__dynamic__ else ""}
241
- """
425
+ if field.type.__dynamic__:
426
+ self._uses_sizes = True
427
+ lines.append("_offset = stream.tell()")
428
+ lines.append(self._assign(field, read))
429
+ lines.append(f"sizes[{field._name!r}] = stream.tell() - _offset")
430
+ else:
431
+ lines.append(self._assign(field, read))
242
432
 
243
- yield dedent(template)
433
+ yield _block(lines)
244
434
 
245
- def _generate_bits(self, field: Field) -> Iterator[str]:
246
- lookup = self._map_field(field)
247
- read_type = "_t"
248
- field_type = field.type
249
- if issubclass(field_type, (Enum, Flag)):
250
- read_type += ".__type__"
251
- field_type = field_type.__type__
435
+ def _generate_field_read_source(self, field: Field, lines: list[str]) -> str:
436
+ """Generate the code for reading a field."""
437
+ context = self._context(lines)
438
+ type_ = self._type_symbol(field.type)
439
+ read = f"{type_}._read(stream, context={context}, endian=endian)"
252
440
 
253
- if issubclass(field_type, Char):
254
- field_type = field_type.__cs__.uint8
255
- lookup = "cls.__cs__.uint8"
441
+ num_entries = getattr(field.type, "__num_entries__", None)
442
+ if not isinstance(num_entries, Expression) or num_entries.expression == "EOF":
443
+ return read
256
444
 
257
- template = f"""
258
- _t = {lookup}
259
- r["{field._name}"] = type.__call__(_t, bit_reader.read({read_type}, {field.bits}))
260
- """
445
+ # If the field is a dynamic array, generate the code for the length evaluation
446
+ # Integer fields that were read before are used directly
447
+ types = {field._name: field.type for field in self.fields}
448
+ variables = {name: f"int({var})" for name, var in self._assigned if issubclass(types[name], _INTEGER_TYPES)}
449
+ try:
450
+ count = num_entries.source(variables, context)
451
+ except ExpressionParserError:
452
+ # Invalid expressions raise their error when reading the array
453
+ return read
454
+
455
+ lines.append(f"_count = max(0, {count})")
456
+ element = self._type_symbol(field.type.__type__)
457
+ return f"_type_call({type_}, {element}._read_array(stream, _count, context={context}, endian=endian))"
458
+
459
+ def _starts_bit_value(self, field: Field) -> bool:
460
+ """Return whether a bit field is read from a new value of the bit buffer."""
461
+ return self._bits is None or self._bits[1] == 0 or self._bits[0] is not _get_bits_type(field)
462
+
463
+ def _generate_bits_source(self, field: Field) -> Iterator[str]:
464
+ """Generate the code for reading a bit field."""
465
+ read_type = _get_bits_type(field)
466
+ new_value = self._starts_bit_value(field)
467
+ if new_value:
468
+ self._bits = (read_type, read_type.__size__ * 8)
469
+ remaining = self._bits[1]
470
+ self._bits = (read_type, remaining - field.bits)
471
+
472
+ lines = []
473
+ if new_value:
474
+ lines.append(f"_bit_buffer = {self._type_symbol(read_type)}._read(stream, endian=endian)")
475
+
476
+ if field.bits > remaining:
477
+ lines.append('raise ValueError("Reading straddled bits is unsupported")')
478
+ yield _block(lines)
479
+ return
480
+
481
+ mask = (1 << field.bits) - 1
482
+ # The bit buffer only reads the least significant bits first for explicit little endian
483
+ lines.append(f'if endian == "<": _bit_value = _bit_buffer & {mask}; _bit_buffer >>= {field.bits}')
484
+ lines.append(f"else: _bit_value = (_bit_buffer >> {remaining - field.bits}) & {mask}")
485
+
486
+ # Character bit fields are returned as the integer type they're read as
487
+ type_ = self._type_symbol(read_type if issubclass(field.type, Char) else field.type)
488
+
489
+ if issubclass(field.type, (Enum, Flag)):
490
+ lines.append(_member_lookup(type_, "_bit_value"))
491
+ lines.append(self._assign(field, "_member"))
492
+ else:
493
+ lines.append(self._assign(field, f"_type_call({type_}, _bit_value)"))
261
494
 
262
- yield dedent(template)
495
+ yield _block(lines)
263
496
 
264
- def _generate_packed(self, fields: list[Field]) -> Iterator[str]:
265
- info = list(_generate_struct_info(self.cs, fields, self.align))
497
+ def _generate_packed_source(self, fields: list[Field]) -> Iterator[str]:
498
+ """Generate the code for reading consecutive packed fields."""
499
+ unpack = _Unpack("data")
266
500
  reads = []
501
+ size = self._pack_fields(fields, 0, unpack, reads, self._assign_local, self.align)
502
+
503
+ lines = [f"buf = stream.read({size})", _check_eof(size)]
504
+ if unpack.has_values():
505
+ lines.append(f"data = {self._add_symbol(unpack.functions('unpack'), 'unpack')}[endian](buf)")
506
+
507
+ yield _block(lines + reads)
508
+
509
+ def _pack_fields(
510
+ self,
511
+ fields: list[Field],
512
+ offset: int,
513
+ unpack: _Unpack,
514
+ lines: list[str],
515
+ target: Callable[[Field, str], list[str]],
516
+ align: bool = False,
517
+ ) -> int:
518
+ """Generate the code for reading the given fields from ``buf``, starting at the given offset.
267
519
 
268
- size = 0
269
- slice_index = 0
270
- for field, count, _ in info:
520
+ Unpackable values are added to the given ``unpack`` format, and accessed through its variable.
521
+ Every field value is passed to ``target``, which returns the lines that store the value.
522
+
523
+ Returns the offset after the last field.
524
+ """
525
+ inline = _inline_kind if self.inline else None
526
+ for field, count, char in _generate_struct_info(fields, align, inline=inline):
271
527
  if field is None:
272
528
  # Padding
273
- size += count
529
+ unpack.add(count, "x")
530
+ offset += count
531
+ continue
532
+
533
+ if char == "S":
534
+ # Inlined structure
535
+ type_ = field.type
536
+ var = self._var("object")
537
+ lines.append(f"{var} = _object_new({self._witness(type_)})")
538
+ end = self._pack_fields(type_.__fields__, offset, unpack, lines, _attr_target(var))
539
+ if padding := offset + type_.__size__ - end:
540
+ unpack.add(padding, "x")
541
+ lines.append(f"{var}.__dynamic_sizes__ = {{}}")
542
+ lines.extend(target(field, var))
543
+ offset += type_.__size__
544
+ continue
545
+
546
+ if char == "U":
547
+ # Inlined union
548
+ type_ = field.type
549
+ var = self._var("union")
550
+ lines.append(f"{var} = _object_new({self._witness(type_)})")
551
+ self._pack_union(var, type_.__fields__, offset, f"buf[{offset}:{offset + type_.__size__}]", lines)
552
+ lines.extend(target(field, var))
553
+ unpack.add(type_.__size__, "x")
554
+ offset += type_.__size__
274
555
  continue
275
556
 
276
557
  field_type = field.type
277
- read_type = _get_read_type(self.cs, field_type)
278
-
279
- if issubclass(field_type, (Array, CharArray, WcharArray)):
280
- count = field_type.__num_entries__
281
- read_type = _get_read_type(self.cs, field_type.__type__)
282
-
283
- if issubclass(read_type, (Char, Wchar, Int)):
284
- count *= read_type.__size__
285
- getter = f"buf[{size}:{size + count}]"
286
- else:
287
- getter = f"data[{slice_index}:{slice_index + count}]"
288
- slice_index += count
289
- elif issubclass(read_type, (Char, Wchar, Int)):
290
- getter = f"buf[{size}:{size + read_type.__size__}]"
558
+ read_type = _get_read_type(field_type)
559
+ is_array = issubclass(field_type, (Array, CharArray, WcharArray))
560
+ if is_array:
561
+ read_type = _get_read_type(field_type.__type__)
562
+
563
+ if char == "x":
564
+ # Byte based types are sliced directly out of the buffer, with the count in bytes
565
+ getter = f"buf[{offset}:{offset + count}]"
566
+ elif is_array:
567
+ getter = f"{unpack.var}[{unpack.index}:{unpack.index + count}]"
291
568
  else:
292
- getter = f"data[{slice_index}]"
293
- slice_index += 1
569
+ getter = f"{unpack.var}[{unpack.index}]"
570
+
571
+ lines.extend(self._generate_packed_field_source(field, read_type, getter, count, target))
572
+ unpack.add(count, char)
573
+ offset += field_type.__size__
294
574
 
295
- if issubclass(read_type, (Wchar, Int)):
296
- # Types that parse bytes further down to their own type
297
- parser_template = "{type}({getter})"
575
+ return offset
576
+
577
+ def _pack_union(self, var: str, fields: list[Field], offset: int, data: str, lines: list[str]) -> None:
578
+ """Generate the code for reading a union from ``buf`` into ``var``.
579
+
580
+ The union starts at the given offset in ``buf``, ``data`` is the code for its own bytes.
581
+ """
582
+ for member in fields:
583
+ member_offset = offset + (member.offset or 0)
584
+ unpack = _Unpack(self._var("data"))
585
+ member_lines = []
586
+ self._pack_fields([member], member_offset, unpack, member_lines, _union_target(var))
587
+
588
+ if unpack.has_values():
589
+ unpack_from = self._add_symbol(unpack.functions("unpack_from"), "unpack_from")
590
+ lines.append(f"{unpack.var} = {unpack_from}[endian](buf, {member_offset})")
591
+ lines.extend(member_lines)
592
+
593
+ lines.append(f"_object_setattr({var}, '__dynamic_sizes__', {{}})")
594
+ lines.append(f"_object_setattr({var}, '_buf', {data})")
595
+ lines.append(f"_object_setattr({var}, '_endian', endian)")
596
+ self._proxify(var, var, fields, lines)
597
+
598
+ def _proxify(self, union: str, var: str, fields: list[Field], lines: list[str]) -> None:
599
+ """Generate the code equivalent to ``Union._proxify``."""
600
+ for field in fields:
601
+ if issubclass(field.type, Structure):
602
+ nested = self._var("nested")
603
+ lines.append(f"{nested} = _getattr({var}, {field._name!r})")
604
+ lines.append(
605
+ f"_object_setattr({var}, {field._name!r}, _UnionProxy({union}, {field._name!r}, {nested}))"
606
+ )
607
+ self._proxify(union, nested, field.type.__fields__, lines)
608
+
609
+ def _generate_packed_field_source(
610
+ self,
611
+ field: Field,
612
+ read_type: type[BaseType],
613
+ getter: str,
614
+ count: int,
615
+ target: Callable[[Field, str], list[str]],
616
+ ) -> list[str]:
617
+ """Generate the code for reading a packed field."""
618
+ field_type = field.type
619
+ lines = []
620
+
621
+ if issubclass(field_type, Array):
622
+ array = self._type_symbol(field.type)
623
+ element = self._type_symbol(field_type.__type__)
624
+
625
+ if issubclass(field_type.__type__, Int):
626
+ elem_size = field_type.__type__.__size__
627
+ signed = bool(field_type.__type__.signed)
628
+ lines.append(f"_array_buf = {getter}")
629
+ raw = f"_int_from_bytes(_array_buf[i:i + {elem_size}], _byteorder[endian], signed={signed})"
630
+ value = f"[_type_call({element}, {raw}) for i in range(0, {count}, {elem_size})]"
631
+ elif issubclass(field_type.__type__, Pointer):
632
+ context = self._context(lines)
633
+ item = f"{element}.__new__({element}, value, stream, context={context}, endian=endian)"
634
+ value = f"[{item} for value in {getter}]"
635
+ elif issubclass(read_type, (Wchar, Int)):
636
+ # Preserve the behaviour of the generic parsing path for these types
637
+ value = f"[{element}(value) for value in {getter}]"
298
638
  else:
299
- # All other types can be simply intialized
300
- parser_template = "type.__call__({type}, {getter})"
301
-
302
- # Create the final reading code
303
- if issubclass(field_type, Array):
304
- reads.append(f"_t = {self._map_field(field)}")
305
- reads.append("_et = _t.__type__")
306
-
307
- if issubclass(field_type.__type__, Int):
308
- reads.append(f"_b = {getter}")
309
- item_parser = parser_template.format(type="_et", getter=f"_b[i:i + {field_type.__type__.__size__}]")
310
- list_comp = f"[{item_parser} for i in range(0, {count}, {field_type.__type__.__size__})]"
311
- elif issubclass(field_type.__type__, Pointer):
312
- item_parser = "_et.__new__(_et, e, stream, context=r, endian=endian)"
313
- list_comp = f"[{item_parser} for e in {getter}]"
314
- else:
315
- item_parser = parser_template.format(type="_et", getter="e")
316
- list_comp = f"[{item_parser} for e in {getter}]"
317
-
318
- parser = f"type.__call__(_t, {list_comp})"
319
- elif issubclass(field_type, CharArray):
320
- parser = f"type.__call__({self._map_field(field)}, {getter})"
321
- elif issubclass(field_type, Pointer):
322
- reads.append(f"_pt = {self._map_field(field)}")
323
- parser = f"_pt.__new__(_pt, {getter}, stream, context=r, endian=endian)"
639
+ value = f"list(map(_type_call, _repeat({element}), {getter}))"
640
+
641
+ lines.extend(target(field, f"_type_call({array}, {value})"))
642
+
643
+ elif issubclass(field_type, CharArray):
644
+ lines.extend(target(field, f"_type_call({self._type_symbol(field.type)}, {getter})"))
645
+
646
+ elif issubclass(field_type, WcharArray) or issubclass(read_type, Wchar):
647
+ encoding = self._add_symbol(read_type.__encoding_map__, "encoding")
648
+ # Using str() instead of .decode() also supports other bytes-like objects returned by the stream
649
+ type_ = self._type_symbol(field.type)
650
+ lines.extend(target(field, f"_type_call({type_}, str({getter}, {encoding}[endian]))"))
651
+
652
+ elif issubclass(field_type, Pointer):
653
+ context = self._context(lines)
654
+ type_ = self._type_symbol(field.type)
655
+ lines.extend(target(field, f"{type_}.__new__({type_}, {getter}, stream, context={context}, endian=endian)"))
656
+
657
+ elif issubclass(field_type, (Enum, Flag)):
658
+ # Fast path the member lookup, the slow path takes care of unknown values
659
+ if issubclass(read_type, Int):
660
+ raw = f"_int_from_bytes({getter}, _byteorder[endian], signed={bool(read_type.signed)})"
324
661
  else:
325
- parser = parser_template.format(type=self._map_field(field), getter=getter)
662
+ raw = getter
326
663
 
327
- reads.append(f'r["{field._name}"] = {parser}')
328
- reads.append("") # Generates a newline in the resulting code
664
+ lines.append(f"_enum_value = {raw}")
665
+ lines.append(_member_lookup(self._type_symbol(field.type), "_enum_value"))
666
+ lines.extend(target(field, "_member"))
329
667
 
330
- size += field_type.__size__
668
+ elif issubclass(read_type, Int):
669
+ signed = bool(read_type.signed)
670
+ value = f"_int_from_bytes({getter}, _byteorder[endian], signed={signed})"
671
+ lines.extend(target(field, f"_type_call({self._type_symbol(field.type)}, {value})"))
331
672
 
332
- fmt = _optimize_struct_fmt(info)
333
- if fmt == "x" or (len(fmt) == 2 and fmt[1] == "x"):
334
- unpack = ""
335
673
  else:
336
- unpack = f'data = _struct(endian, "{fmt}").unpack(buf)\n'
674
+ lines.extend(target(field, f"_type_call({self._type_symbol(field.type)}, {getter})"))
337
675
 
338
- template = f"""
339
- buf = stream.read({size})
340
- if len(buf) != {size}: raise EOFError()
341
- {unpack}
342
- """
676
+ return lines
677
+
678
+
679
+ class _Unpack:
680
+ """Tracks consecutive struct unpack formats."""
681
+
682
+ def __init__(self, var: str):
683
+ self.var = var
684
+ self.index = 0
685
+ self.items: list[tuple[int, str]] = []
686
+
687
+ def add(self, count: int, char: str) -> None:
688
+ self.items.append((count, char))
689
+ if char != "x":
690
+ self.index += count
691
+
692
+ def has_values(self) -> bool:
693
+ return self.index > 0
694
+
695
+ def functions(self, name: str) -> dict[str, Callable]:
696
+ """Return the given function of the struct for every endianness."""
697
+ fmt = _optimize_struct_fmt(self.items)
698
+ return {endian: getattr(_struct(endian, fmt), name) for endian in _ENDIANNESS}
699
+
700
+
701
+ def _check_union(cls: type[Union], fields: list[Field], inline: bool = True) -> int:
702
+ """Check if the compiler supports the union. Returns its size."""
703
+ size = cls.__size__
704
+ if (
705
+ not inline
706
+ or cls.__align__
707
+ or size is None
708
+ or not _uses_generated_init(cls)
709
+ or not _can_inline_fields(fields, size, is_union=True)
710
+ ):
711
+ raise TypeError(f"Unsupported union for compiler: {cls}")
712
+ return size
713
+
714
+
715
+ def _inline_kind(type_: type[BaseType], in_union: bool = False) -> str | None:
716
+ """Determine if a structure type can be inlined, returning ``S`` for structures and ``U`` for unions."""
717
+ if not issubclass(type_, Structure):
718
+ return None
719
+
720
+ # Only static structures that we can construct directly, the same way their (compiled) _read would
721
+ if type_.__size__ is None or type_.__align__ or not _uses_generated_init(type_):
722
+ return None
723
+
724
+ is_union = issubclass(type_, Union)
725
+ if (is_union and in_union) or (not is_union and not type_.__compiled__):
726
+ # Unions nested in unions aren't supported when proxying
727
+ return None
728
+
729
+ if not _can_inline_fields(type_.__fields__, type_.__size__, is_union, in_union):
730
+ return None
731
+
732
+ return "U" if is_union else "S"
733
+
734
+
735
+ def _can_inline_fields(fields: list[Field], size: int, is_union: bool, in_union: bool = False) -> bool:
736
+ """Determine if the fields of a structure or union type of the given size can be inlined."""
737
+ end = 0
738
+ for field in fields:
739
+ # Union members without an offset are read from the start of the union
740
+ if field.bits or (field.offset is None and not is_union) or (not is_union and field.offset < end):
741
+ return False
742
+
743
+ field_type = field.type
744
+ if isinstance(field_type, EnumMetaType):
745
+ field_type = field_type.__type__
746
+
747
+ if not issubclass(field_type, SUPPORTED_TYPES) or issubclass(field_type, (Pointer, Void, VoidArray)):
748
+ return False
749
+
750
+ if issubclass(field_type, Structure):
751
+ if _inline_kind(field_type, in_union or is_union) is None:
752
+ return False
753
+ elif issubclass(field_type, (Array, CharArray, WcharArray)) and (
754
+ field_type.__dynamic__ or issubclass(field_type.__type__, (Structure, BaseArray, Pointer, Void))
755
+ ):
756
+ return False
757
+
758
+ # Union members can be placed beyond the end of the union, which reads past its bytes
759
+ if (end := (field.offset or 0) + field_type.__size__) > size:
760
+ return False
761
+
762
+ return True
763
+
764
+
765
+ def _uses_generated_init(type_: type[Structure]) -> bool:
766
+ """Return whether a structure type uses a generated ``__init__``, so we can bypass it by setting attributes."""
767
+ return "__init__" in type_.__dict__ and type_.__new__ is object.__new__
768
+
769
+
770
+ def _set_attribute(var: str, name: str, value: str) -> str:
771
+ """Return the code to set an attribute."""
772
+ if name.isidentifier() and not iskeyword(name):
773
+ return f"{var}.{name} = {value}"
774
+ return f"_setattr({var}, {name!r}, {value})"
775
+
776
+
777
+ def _attr_target(var: str) -> Callable[[Field, str], list[str]]:
778
+ """Return a target that sets field values as attributes on the given object variable."""
779
+
780
+ def target(field: Field, value: str) -> list[str]:
781
+ return [_set_attribute(var, field._name, value)]
343
782
 
344
- yield dedent(template) + "\n".join(reads)
783
+ return target
345
784
 
346
785
 
347
- def _generate_struct_info(cs: cstruct, fields: list[Field], align: bool = False) -> Iterator[tuple[Field, int, str]]:
786
+ def _union_target(var: str) -> Callable[[Field, str], list[str]]:
787
+ """Return a target that sets field values on the given union object variable."""
788
+
789
+ def target(field: Field, value: str) -> list[str]:
790
+ # This bypasses `__setattr__`
791
+ return [f"_object_setattr({var}, {field._name!r}, {value})"]
792
+
793
+ return target
794
+
795
+
796
+ def _member_lookup(type_: str, value: str) -> str:
797
+ """Return the code to get the enum or flag member for a value into ``_member``."""
798
+ return f"if (_member := {type_}._value2member_map_.get({value})) is None: _member = _type_call({type_}, {value})"
799
+
800
+
801
+ def _check_eof(size: int) -> str:
802
+ """Return the code to raise an ``EOFError`` if fewer than ``size`` bytes were read into ``buf``."""
803
+ return f'if len(buf) != {size}: raise EOFError(f"Read {{len(buf)}} bytes, but expected {size}")'
804
+
805
+
806
+ def _block(lines: list[str]) -> str:
807
+ return "\n" + "\n".join(lines) + "\n"
808
+
809
+
810
+ def _generate_struct_info(
811
+ fields: list[Field],
812
+ align: bool = False,
813
+ inline: Callable[[type[BaseType]], str | None] | None = None,
814
+ ) -> Iterator[tuple[Field, int, str]]:
348
815
  if not fields:
349
816
  return
350
817
 
@@ -362,16 +829,20 @@ def _generate_struct_info(cs: cstruct, fields: list[Field], align: bool = False)
362
829
  imaginary_offset += drift
363
830
 
364
831
  count = 1
365
- read_type = _get_read_type(cs, field.type)
832
+ read_type = _get_read_type(field.type)
366
833
 
367
834
  # Drop voids
368
835
  if issubclass(read_type, (Void, VoidArray)):
369
836
  continue
370
837
 
838
+ # Structures that are inlined, yield as (S)tructure or (U)nion
839
+ if inline is not None and (kind := inline(read_type)) is not None:
840
+ yield field, 1, kind
841
+
371
842
  # Array of more complex types are handled elsewhere
372
843
  if issubclass(read_type, (Array, CharArray, WcharArray)):
373
844
  count = read_type.__num_entries__
374
- read_type = _get_read_type(cs, read_type.__type__)
845
+ read_type = _get_read_type(read_type.__type__)
375
846
 
376
847
  # Take the pack char for Packed
377
848
  if issubclass(read_type, Packed):
@@ -388,37 +859,28 @@ def _generate_struct_info(cs: cstruct, fields: list[Field], align: bool = False)
388
859
  current_offset += size
389
860
 
390
861
 
391
- def _optimize_struct_fmt(info: Iterator[tuple[Field, int, str]]) -> str:
862
+ def _optimize_struct_fmt(items: Iterable[tuple[int, str]]) -> str:
863
+ """Create a struct format string from ``(count, char)`` items, merging consecutive items of the same char."""
392
864
  chars = []
865
+ for char, group in groupby(items, key=itemgetter(1)):
866
+ if count := sum(count for count, _ in group):
867
+ chars.append(f"{count if count > 1 else ''}{char}")
393
868
 
394
- current_count = 0
395
- current_char = None
396
-
397
- for _, count, char in info:
398
- if current_char is None:
399
- current_count = count
400
- current_char = char
401
- continue
402
-
403
- if char != current_char:
404
- if current_count:
405
- chars.append((current_count, current_char))
406
- current_count = count
407
- current_char = char
408
- else:
409
- current_count += count
869
+ return "".join(chars)
410
870
 
411
- if current_char is not None and current_count:
412
- chars.append((current_count, current_char))
413
871
 
414
- return "".join(f"{count if count > 1 else ''}{char}" for count, char in chars)
872
+ def _get_bits_type(field: Field) -> type[BaseType]:
873
+ """Return the type that ``BitBuffer`` reads a bit field as."""
874
+ type_ = field.type.__type__ if issubclass(field.type, (Enum, Flag)) else field.type
875
+ return type_.__cs__.uint8 if issubclass(type_, Char) else type_
415
876
 
416
877
 
417
- def _get_read_type(cs: cstruct, type_: type[BaseType]) -> type[BaseType]:
878
+ def _get_read_type(type_: type[BaseType]) -> type[BaseType]:
879
+ """Return the type that a value of the given type is read as."""
418
880
  if issubclass(type_, (Enum, Flag)):
419
881
  type_ = type_.__type__
420
882
 
421
883
  if issubclass(type_, Pointer):
422
- type_ = cs.pointer
884
+ type_ = type_.__cs__.pointer
423
885
 
424
- return cs.resolve(type_)
886
+ return type_