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

dissect/cstruct/compiler.py

@@ -182,8 +182,10 @@ _GLOBALS = {
 }
 
 
-class _ReadSourceGenerator:
-    """Generates the source code of an optimized ``_read`` method for a structure or union."""
+class _SourceGenerator:
+    """Base class for generating the source code of an optimized method of a structure or union."""
+
+    name: str
 
     def __init__(self, cls: type[Structure], fields: list[Field], inline: bool = True):
         self.cls = cls
@@ -191,24 +193,17 @@ class _ReadSourceGenerator:
         self.align = cls.__align__
         self.inline = inline and not self.align
 
+        # The values available to the generated code
+        self.symbols: dict[str, Any] = {"cs": cls.__cs__}
+        self._symbol_names: dict[int, str] = {}
+
         # Witnesses of inlined structure types (symbol of the type -> symbol of its __init__ at compile time)
         self._witnesses: dict[str, str] = {}
-        self._fallback_symbol: str | None = None
 
         # The state of the bit buffer: the type of its current value and the remaining bits
         self._bits: tuple[type[BaseType], int] | None = None
         self._var_ids: dict[str, int] = {}
 
-        # The values available to the generated code
-        self.symbols: dict[str, Any] = {"cs": cls.__cs__}
-        self._symbol_names: dict[int, str] = {}
-
-        # (field name, local variable name) of all fields assigned so far
-        self._assigned: list[tuple[str, str]] = []
-        self._has_context = False
-        self._uses_start = False
-        self._uses_sizes = False
-
     def _type_symbol(self, type_: type[BaseType]) -> str:
         """Add a type as symbol, returns its name."""
         return self._add_symbol(type_, "type")
@@ -220,39 +215,12 @@ class _ReadSourceGenerator:
             self.symbols[symbol] = value
         return symbol
 
-    def _context(self, lines: list[str]) -> str:
-        """Return the name of the context dictionary, creating it from the fields read so far if needed."""
-        if not self._has_context:
-            self._has_context = True
-            items = ", ".join(f"{name!r}: {var}" for name, var in self._assigned)
-            lines.append(f"values = {{{items}}}")
-
-        return "values"
-
     def _var(self, name: str) -> str:
         """Return a unique variable name for the generated code."""
         index = self._var_ids.get(name, 0)
         self._var_ids[name] = index + 1
         return f"_{name}{index}"
 
-    def _assign(self, field: Field, value: str) -> str:
-        """Assign a value to a new local variable for the given field, and the context dictionary if it exists."""
-        var = f"_value{len(self._assigned)}"
-        self._assigned.append((field._name, var))
-        if self._has_context:
-            return f"{var} = values[{field._name!r}] = {value}"
-        return f"{var} = {value}"
-
-    def _assign_local(self, field: Field, value: str) -> list[str]:
-        return [self._assign(field, value)]
-
-    def _fallback(self) -> str:
-        """Return the symbol of the non-inlined fallback read function."""
-        if self._fallback_symbol is None:
-            fallback = _Fallback(self.cls, self.fields)
-            self._fallback_symbol = self._add_symbol(fallback, "fallback")
-        return self._fallback_symbol
-
     def _witness(self, type_: type[Structure]) -> str:
         """Add an inlined type as symbol, with a witness to detect modifications to it after compilation."""
         symbol = self._type_symbol(type_)
@@ -263,10 +231,22 @@ class _ReadSourceGenerator:
         """Return the conditions that are true if any of the inlined types were modified after compilation."""
         return [f"{type_}.__init__ is not {init}" for type_, init in self._witnesses.items()]
 
+    def _starts_bit_value(self, field: Field) -> bool:
+        """Return whether a bit field is the first of a new value of the bit buffer."""
+        return self._bits is None or self._bits[1] == 0 or self._bits[0] is not _get_bit_buffer_type(field)
+
+    def _advance_bits(self, field: Field) -> int:
+        """Update the state of the bit buffer for a bit field, returns the bits that were remaining before it."""
+        if self._starts_bit_value(field):
+            self._bits = (_get_bit_buffer_type(field), _get_bits_type(field).__size__ * 8)
+        remaining = self._bits[1]
+        self._bits = (self._bits[0], remaining - field.bits)
+        return remaining
+
     def generate(self) -> MethodType:
         source = self.generate_source()
 
-        code = python_compile(source, f"<compiled {self.cls.__name__}._read>", "exec")
+        code = python_compile(source, f"<compiled {self.cls.__name__}.{self.name}>", "exec")
         exec(code, {**_GLOBALS, **self.symbols}, d := {})
         obj = d.popitem()[1]
         obj.__source__ = source
@@ -278,6 +258,56 @@ class _ReadSourceGenerator:
             return self._generate_union_source()
         return self._generate_struct_source()
 
+    def _generate_struct_source(self) -> str:
+        raise NotImplementedError
+
+    def _generate_union_source(self) -> str:
+        raise NotImplementedError
+
+
+class _ReadSourceGenerator(_SourceGenerator):
+    """Generates the source code of an optimized ``_read`` method for a structure or union."""
+
+    name = "_read"
+
+    def __init__(self, cls: type[Structure], fields: list[Field], inline: bool = True):
+        super().__init__(cls, fields, inline)
+
+        self._fallback_symbol: str | None = None
+
+        # (field name, local variable name) of all fields assigned so far
+        self._assigned: list[tuple[str, str]] = []
+        self._has_context = False
+        self._uses_start = False
+        self._uses_sizes = False
+
+    def _context(self, lines: list[str]) -> str:
+        """Return the name of the context dictionary, creating it from the fields read so far if needed."""
+        if not self._has_context:
+            self._has_context = True
+            items = ", ".join(f"{name!r}: {var}" for name, var in self._assigned)
+            lines.append(f"values = {{{items}}}")
+
+        return "values"
+
+    def _assign(self, field: Field, value: str) -> str:
+        """Assign a value to a new local variable for the given field, and the context dictionary if it exists."""
+        var = f"_value{len(self._assigned)}"
+        self._assigned.append((field._name, var))
+        if self._has_context:
+            return f"{var} = values[{field._name!r}] = {value}"
+        return f"{var} = {value}"
+
+    def _assign_local(self, field: Field, value: str) -> list[str]:
+        return [self._assign(field, value)]
+
+    def _fallback(self) -> str:
+        """Return the symbol of the non-inlined fallback read function."""
+        if self._fallback_symbol is None:
+            fallback = _Fallback(self.cls, self.fields)
+            self._fallback_symbol = self._add_symbol(fallback, "fallback")
+        return self._fallback_symbol
+
     def _generate_struct_source(self) -> str:
         """Generate the code for reading a structure."""
         read_code = "\n".join(self._generate_fields_source())
@@ -405,6 +435,16 @@ class _ReadSourceGenerator:
 
             # Everything else - basic and composite types (and arrays of them), and inlined structures
             else:
+                if field.offset is None:
+                    # The alignment depends on the position in the stream, so this needs a new block
+                    if self.align:
+                        yield from flush()
+                # A block is read in order, so a field before the current position needs a new block
+                elif current_offset is None or field.offset < current_offset:
+                    yield from flush()
+                elif current_block:
+                    current_offset = field.offset
+
                 if not current_block:
                     yield from align_to_field(field)
                 current_block.append(field)
@@ -456,18 +496,11 @@ class _ReadSourceGenerator:
         element = self._type_symbol(field.type.__type__)
         return f"_type_call({type_}, {element}._read_array(stream, _count, context={context}, endian=endian))"
 
-    def _starts_bit_value(self, field: Field) -> bool:
-        """Return whether a bit field is read from a new value of the bit buffer."""
-        return self._bits is None or self._bits[1] == 0 or self._bits[0] is not _get_bits_type(field)
-
     def _generate_bits_source(self, field: Field) -> Iterator[str]:
         """Generate the code for reading a bit field."""
         read_type = _get_bits_type(field)
         new_value = self._starts_bit_value(field)
-        if new_value:
-            self._bits = (read_type, read_type.__size__ * 8)
-        remaining = self._bits[1]
-        self._bits = (read_type, remaining - field.bits)
+        remaining = self._advance_bits(field)
 
         lines = []
         if new_value:
@@ -498,7 +531,7 @@ class _ReadSourceGenerator:
         """Generate the code for reading consecutive packed fields."""
         unpack = _Unpack("data")
         reads = []
-        size = self._pack_fields(fields, 0, unpack, reads, self._assign_local, self.align)
+        size = self._pack_fields(fields, 0, unpack, reads, self._assign_local)
 
         lines = [f"buf = stream.read({size})", _check_eof(size)]
         if unpack.has_values():
@@ -513,7 +546,6 @@ class _ReadSourceGenerator:
         unpack: _Unpack,
         lines: list[str],
         target: Callable[[Field, str], list[str]],
-        align: bool = False,
     ) -> int:
         """Generate the code for reading the given fields from ``buf``, starting at the given offset.
 
@@ -523,7 +555,7 @@ class _ReadSourceGenerator:
         Returns the offset after the last field.
         """
         inline = _inline_kind if self.inline else None
-        for field, count, char in _generate_struct_info(fields, align, inline=inline):
+        for field, count, char in _generate_struct_info(fields, inline=inline):
             if field is None:
                 # Padding
                 unpack.add(count, "x")
@@ -595,16 +627,15 @@ class _ReadSourceGenerator:
         lines.append(f"_object_setattr({var}, '_endian', endian)")
         self._proxify(var, var, fields, lines)
 
-    def _proxify(self, union: str, var: str, fields: list[Field], lines: list[str]) -> None:
+    def _proxify(self, union: str, var: str, fields: list[Field], lines: list[str], attr: str | None = None) -> None:
         """Generate the code equivalent to ``Union._proxify``."""
         for field in fields:
             if issubclass(field.type, Structure):
                 nested = self._var("nested")
+                member = attr or field._name
                 lines.append(f"{nested} = _getattr({var}, {field._name!r})")
-                lines.append(
-                    f"_object_setattr({var}, {field._name!r}, _UnionProxy({union}, {field._name!r}, {nested}))"
-                )
-                self._proxify(union, nested, field.type.__fields__, lines)
+                lines.append(f"_object_setattr({var}, {field._name!r}, _UnionProxy({union}, {member!r}, {nested}))")
+                self._proxify(union, nested, field.type.__fields__, lines, member)
 
     def _generate_packed_field_source(
         self,
@@ -622,19 +653,18 @@ class _ReadSourceGenerator:
             array = self._type_symbol(field.type)
             element = self._type_symbol(field_type.__type__)
 
-            if issubclass(field_type.__type__, Int):
-                elem_size = field_type.__type__.__size__
-                signed = bool(field_type.__type__.signed)
+            if issubclass(read_type, Int):
+                # Everything that's read as an integer type is sliced directly out of the buffer
+                elem_size = read_type.__size__
+                signed = bool(read_type.signed)
                 lines.append(f"_array_buf = {getter}")
                 raw = f"_int_from_bytes(_array_buf[i:i + {elem_size}], _byteorder[endian], signed={signed})"
-                value = f"[_type_call({element}, {raw}) for i in range(0, {count}, {elem_size})]"
-            elif issubclass(field_type.__type__, Pointer):
+                getter = f"[{raw} for i in range(0, {count}, {elem_size})]"
+
+            if issubclass(field_type.__type__, Pointer):
                 context = self._context(lines)
                 item = f"{element}.__new__({element}, value, stream, context={context}, endian=endian)"
                 value = f"[{item} for value in {getter}]"
-            elif issubclass(read_type, (Wchar, Int)):
-                # Preserve the behaviour of the generic parsing path for these types
-                value = f"[{element}(value) for value in {getter}]"
             else:
                 value = f"list(map(_type_call, _repeat({element}), {getter}))"
 
@@ -647,7 +677,7 @@ class _ReadSourceGenerator:
             encoding = self._add_symbol(read_type.__encoding_map__, "encoding")
             # Using str() instead of .decode() also supports other bytes-like objects returned by the stream
             type_ = self._type_symbol(field.type)
-            lines.extend(target(field, f"_type_call({type_}, str({getter}, {encoding}[endian]))"))
+            lines.extend(target(field, f"_type_call({type_}, str({getter}, {encoding}[endian], 'surrogatepass'))"))
 
         elif issubclass(field_type, Pointer):
             context = self._context(lines)
@@ -694,8 +724,7 @@ class _Unpack:
 
     def functions(self, name: str) -> dict[str, Callable]:
         """Return the given function of the struct for every endianness."""
-        fmt = _optimize_struct_fmt(self.items)
-        return {endian: getattr(_struct(endian, fmt), name) for endian in _ENDIANNESS}
+        return _struct_functions(_optimize_struct_fmt(self.items), name)
 
 
 def _check_union(cls: type[Union], fields: list[Field], inline: bool = True) -> int:
@@ -809,25 +838,18 @@ def _block(lines: list[str]) -> str:
 
 def _generate_struct_info(
     fields: list[Field],
-    align: bool = False,
     inline: Callable[[type[BaseType]], str | None] | None = None,
 ) -> Iterator[tuple[Field, int, str]]:
     if not fields:
         return
 
     current_offset = fields[0].offset
-    imaginary_offset = 0
     for field in fields:
         # We moved -- probably due to alignment
         if field.offset is not None and (drift := field.offset - current_offset) > 0:
             yield None, drift, "x"
             current_offset += drift
 
-        if align and field.offset is None and (drift := -imaginary_offset & (field.alignment - 1)) > 0:
-            # Assume we started at a correctly aligned boundary
-            yield None, drift, "x"
-            imaginary_offset += drift
-
         count = 1
         read_type = _get_read_type(field.type)
 
@@ -853,25 +875,41 @@ def _generate_struct_info(
         elif issubclass(read_type, (Char, Wchar, Int)):
             yield field, count * read_type.__size__, "x"
 
-        size = count * read_type.__size__
-        imaginary_offset += size
         if current_offset is not None:
-            current_offset += size
+            current_offset += count * read_type.__size__
 
 
 def _optimize_struct_fmt(items: Iterable[tuple[int, str]]) -> str:
-    """Create a struct format string from ``(count, char)`` items, merging consecutive items of the same char."""
+    """Create a struct format string from ``(count, char)`` items, merging consecutive items of the same char.
+
+    Byte strings aren't merged, since those are a single value regardless of their count.
+    """
     chars = []
     for char, group in groupby(items, key=itemgetter(1)):
-        if count := sum(count for count, _ in group):
+        if char == "s":
+            chars.extend(f"{count}{char}" for count, _ in group)
+        elif count := sum(count for count, _ in group):
             chars.append(f"{count if count > 1 else ''}{char}")
 
     return "".join(chars)
 
 
+def _struct_functions(fmt: str, name: str) -> dict[str, Callable]:
+    """Return the given function of a struct with the given format for every endianness.
+
+    The ``@`` endianness uses the ``=`` format, since ``@`` would add padding between items to align them.
+    """
+    return {endian: getattr(_struct("=" if endian == "@" else endian, fmt), name) for endian in _ENDIANNESS}
+
+
+def _get_bit_buffer_type(field: Field) -> type[BaseType]:
+    """Return the type that ``BitBuffer`` uses for a bit field, only bit fields of the same type share a value."""
+    return field.type.__type__ if issubclass(field.type, (Enum, Flag)) else field.type
+
+
 def _get_bits_type(field: Field) -> type[BaseType]:
     """Return the type that ``BitBuffer`` reads a bit field as."""
-    type_ = field.type.__type__ if issubclass(field.type, (Enum, Flag)) else field.type
+    type_ = _get_bit_buffer_type(field)
     return type_.__cs__.uint8 if issubclass(type_, Char) else type_
 
 

dissect/cstruct/cstruct.py

@@ -86,7 +86,7 @@ class cstruct:
             "int24": self._make_int_type("int24", 3, True, alignment=4),
             "uint24": self._make_int_type("uint24", 3, False, alignment=4),
             "int48": self._make_int_type("int48", 6, True, alignment=8),
-            "uint48": self._make_int_type("int48", 6, False, alignment=8),
+            "uint48": self._make_int_type("uint48", 6, False, alignment=8),
             "int128": self._make_int_type("int128", 16, True, alignment=16),
             "uint128": self._make_int_type("uint128", 16, False, alignment=16),
 

dissect/cstruct/types/base.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import functools
 import warnings
 from io import SEEK_CUR, BytesIO
+from types import MethodType
 from typing import TYPE_CHECKING, Any, BinaryIO, ClassVar, TypeVar
 
 from dissect.cstruct.exception import ArraySizeError
@@ -19,6 +20,8 @@ if TYPE_CHECKING:
 
 EOF = -0xE0F  # Negative counts are illegal anyway, so abuse that for our EOF sentinel
 
+_IMMUTABLE_TYPES = (int, float, bytes, str)
+
 
 class MetaType(type):
     """Base metaclass for cstruct type classes."""
@@ -218,7 +221,8 @@ class MetaType(type):
             array: The array to write.
             endian: The endianness to use when parsing (little, big, network, <, >, !, @ or =).
         """
-        return sum(cls._write(stream, entry, endian=endian) for entry in array)
+        write = cls._write
+        return sum(write(stream, entry, endian=endian) for entry in array)
 
     def _write_0(cls, stream: BinaryIO, array: list[Self], *, endian: Endianness) -> int:  # type: ignore
         """Internal function for writing null-terminated arrays.
@@ -295,20 +299,29 @@ class _overload:
         b'\\x7b\\x00\\x00\\x00'
     """
 
-    def __init__(self, func: Callable[..., Any]) -> None:
+    def __init__(self, func: Callable[..., Any], method: Callable[..., Any]) -> None:
         self.func = func
+        self.method = method
 
-    def __get__(self, instance: BaseType | None, owner: type[BaseType]) -> Callable[[], bytes]:
+    def __get__(self, instance: BaseType | None, owner: type[BaseType]) -> Callable[..., Any]:
         if instance is None:
             return functools.partial(self.func, owner)
-        return functools.partial(self.func, instance.__class__, value=instance)
+        return MethodType(self.method, instance)
+
+
+def _instance_dumps(self: BaseType, *, endian: AllowedEndianness | None = None) -> bytes:
+    return MetaType.dumps(self.__class__, self, endian=endian)
+
+
+def _instance_write(self: BaseType, stream: BinaryIO, *, endian: AllowedEndianness | None = None) -> int:
+    return MetaType.write(self.__class__, stream, self, endian=endian)
 
 
 class BaseType(metaclass=MetaType):
     """Base class for cstruct type classes."""
 
-    dumps = _overload(MetaType.dumps)
-    write = _overload(MetaType.write)
+    dumps = _overload(MetaType.dumps, _instance_dumps)
+    write = _overload(MetaType.write, _instance_write)
 
     def __len__(self) -> int:
         """Return the byte size of the type."""
@@ -338,9 +351,10 @@ class BaseArray(BaseType):
 
     @classmethod
     def __default__(cls) -> BaseType:
-        return type.__call__(
-            cls, [cls.__type__.__default__()] * (cls.__num_entries__ if isinstance(cls.__num_entries__, int) else 0)
-        )
+        count = cls.__num_entries__ if isinstance(cls.__num_entries__, int) else 0
+        if isinstance(value := cls.__type__.__default__(), _IMMUTABLE_TYPES):
+            return type.__call__(cls, [value] * count)
+        return type.__call__(cls, [cls.__type__.__default__() for _ in range(count)])
 
     @classmethod
     def _read(cls, stream: BinaryIO, *, context: dict[str, Any] | None = None, endian: Endianness) -> list[BaseType]:

dissect/cstruct/types/char.py

@@ -23,8 +23,9 @@ class CharArray(bytes, BaseArray):
         return type.__call__(cls, super()._read(stream, context=context, endian=endian))
 
     @classmethod
-    def _write(cls, stream: BinaryIO, data: bytes | str | list[int], *, endian: Endianness) -> int:
-        if isinstance(data, list):
+    def _write(cls, stream: BinaryIO, data: bytes | bytearray | str | list[int], *, endian: Endianness) -> int:
+        # Copy mutable data types to bytes so we don't modify them in place
+        if isinstance(data, (list, bytearray)):
             data = bytes(data)
         elif isinstance(data, str):
             data = data.encode("latin-1")
@@ -92,4 +93,7 @@ class Char(bytes, BaseType):
         if isinstance(data, str):
             data = data.encode("latin-1")
 
-        return stream.write(data)
+        if (actual_size := len(data)) > 1:
+            raise ValueError(f"Expected a single character, got {actual_size} instead.")
+
+        return stream.write(data or b"\x00")

dissect/cstruct/types/structure.py

@@ -7,7 +7,7 @@ from collections.abc import MutableMapping
 from contextlib import contextmanager
 from enum import Enum
 from functools import lru_cache
-from itertools import chain
+from itertools import chain, repeat
 from operator import attrgetter
 from textwrap import dedent
 from types import FunctionType
@@ -18,6 +18,7 @@ from dissect.cstruct.expression import Expression
 from dissect.cstruct.lexer import IDENTIFIER_TYPES
 from dissect.cstruct.types.base import (
     _BUFFER_TYPES,
+    _IMMUTABLE_TYPES,
     BaseArray,
     BaseType,
     MetaType,
@@ -141,7 +142,6 @@ class StructureMetaType(MetaType):
         else:
             size, alignment = cls._calculate_size_and_offsets(fields, align)
 
-        # TODO: compile _write
         # TODO: generate cached_property for lazy reading
 
         classdict["__size__"] = size
@@ -193,7 +193,7 @@ class StructureMetaType(MetaType):
                     # Moved to a bit field of another type, e.g. uint16 f1 : 8, uint32 f2 : 8;
                     or field_type != bits_type
                     # Still processing a bit field, but it's at a different offset due to alignment or a manual offset
-                    or (bits_type is not None and offset > bits_field_offset + bits_type.__size__)
+                    or (None not in (offset, bits_field_offset) and offset > bits_field_offset + bits_type.__size__)
                 ):
                     # ... if any of this is true, we have to move to the next field
                     bits_type = field_type
@@ -299,7 +299,6 @@ class StructureMetaType(MetaType):
     def _write(cls, stream: BinaryIO, data: Structure, *, endian: Endianness) -> int:
         bit_buffer = None
         struct_start = stream.tell()
-        num = 0
 
         for field in cls.__fields__:
             field_type = field.type
@@ -331,7 +330,7 @@ class StructureMetaType(MetaType):
                     stream.write(b"\x00" * align_pad)
                     offset += align_pad
 
-            value = getattr(data, field._name, None)
+            value = getattr(data, field._name)
             if value is None:
                 value = field_type.__default__()
 
@@ -345,7 +344,6 @@ class StructureMetaType(MetaType):
                     bit_buffer.write(field_type, value, field.bits)
             else:
                 field_type._write(stream, value, endian=endian)
-                num += stream.tell() - offset
 
         if bit_buffer is not None and bit_buffer._type is not None:
             bit_buffer.flush()
@@ -354,7 +352,7 @@ class StructureMetaType(MetaType):
             # Align the stream
             stream.write(b"\x00" * (-stream.tell() & (cls.__alignment__ - 1)))
 
-        return num
+        return stream.tell() - struct_start
 
     def add_field(cls, name: str, type_: type[BaseType], bits: int | None = None, offset: int | None = None) -> None:
         field = Field(name, type_, bits=bits, offset=offset)
@@ -698,6 +696,10 @@ class UnionMetaType(StructureMetaType):
 
         return obj
 
+    def _write_field(cls) -> Field | None:
+        """Return the field to write a union by, its largest field and preferably not an anonymous structure."""
+        return max(cls.__fields__, key=lambda e: (e.type.__size__ or 0, e.name is not None), default=None)
+
     def _write(cls, stream: BinaryIO, data: Union, *, endian: Endianness) -> int:
         if cls.__dynamic__:
             raise NotImplementedError("Writing dynamic unions is not yet supported")
@@ -705,24 +707,9 @@ class UnionMetaType(StructureMetaType):
         offset = stream.tell()
         expected_offset = offset + len(cls)
 
-        # Sort by largest field
-        fields = sorted(cls.__fields__, key=lambda e: e.type.__size__ or 0, reverse=True)
-        anonymous_struct = False
-
-        # Try to write by largest field
-        for field in fields:
-            if isinstance(field.type, StructureMetaType) and field.name is None:
-                # Prefer to write regular fields initially
-                anonymous_struct = field.type
-                continue
-
+        if (field := cls._write_field()) is not None:
             # Write the value
             field.type._write(stream, getattr(data, field._name), endian=endian)
-            break
-
-        # If we haven't written anything yet and we initially skipped an anonymous struct, write it now
-        if stream.tell() == offset and anonymous_struct:
-            anonymous_struct._write(stream, data, endian=endian)
 
         # If we haven't filled the union size yet, pad it
         if remaining := expected_offset - stream.tell():
@@ -745,7 +732,9 @@ class Union(Structure, metaclass=UnionMetaType):
             raise NotImplementedError("Modifying a dynamic union is not yet supported")
 
         super().__setattr__(attr, value)
-        self._rebuild(attr)
+        # Fields of an anonymous structure are rebuilt through the proxy of that structure
+        if attr in self.__class__.__lookup__:
+            self._rebuild(attr)
 
     def _rebuild(self, attr: str) -> None:
         if (cur_buf := getattr(self, "_buf", None)) is None:
@@ -773,13 +762,17 @@ class Union(Structure, metaclass=UnionMetaType):
         object.__setattr__(self, "__dynamic_sizes__", sizes)
 
     def _proxify(self) -> None:
-        def _proxy_structure(value: Structure) -> None:
+        def _proxy_structure(value: Structure, attr: str | None = None) -> None:
             for field in value.__class__.__fields__:
                 if issubclass(field.type, Structure):
                     nested_value = getattr(value, field._name)
-                    proxy = UnionProxy(self, field._name, nested_value)
-                    object.__setattr__(value, field._name, proxy)
-                    _proxy_structure(nested_value)
+                    # Changing a nested field rebuilds the union from the top level field it's in
+                    member = attr or field._name
+                    object.__setattr__(value, field._name, UnionProxy(self, member, nested_value))
+                    # The value of a nested union is already a proxy, so continue with the value behind it
+                    if isinstance(nested_value, UnionProxy):
+                        nested_value = nested_value.__target__
+                    _proxy_structure(nested_value, member)
 
         _proxy_structure(self)
 
@@ -864,7 +857,7 @@ def _make_structure__init__(fields: list[str]) -> str:
     """
     field_args = ", ".join(f"{field} = None" for field in fields)
     field_init = "\n".join(
-        f" self.{name} = {name} if {name} is not None else _{i}_default" for i, name in enumerate(fields)
+        f" self.{name} = {name} if {name} is not None else _{i}_default()" for i, name in enumerate(fields)
     )
 
     code = f"def __init__(self{', ' + field_args or ''}):\n"
@@ -880,7 +873,7 @@ def _make_union__init__(fields: list[str]) -> str:
     """
     field_args = ", ".join(f"{field} = None" for field in fields)
     field_init = "\n".join(
-        f" object.__setattr__(self, '{name}', {name} if {name} is not None else _{i}_default)"
+        f" object.__setattr__(self, '{name}', {name} if {name} is not None else _{i}_default())"
         for i, name in enumerate(fields)
     )
 
@@ -976,7 +969,7 @@ def _generate_structure__init__(fields: list[Field]) -> FunctionType:
             co_names=_remap_co_values(template.__code__.co_names, mapping),
             co_varnames=_remap_co_values(template.__code__.co_varnames, mapping),
         ),
-        template.__globals__ | {f"__{field._name}_default__": field.type.__default__() for field in fields},
+        template.__globals__ | {f"__{field._name}_default__": _default_factory(field.type) for field in fields},
         argdefs=template.__defaults__,
     )
 
@@ -996,11 +989,17 @@ def _generate_union__init__(fields: list[Field]) -> FunctionType:
             co_names=_remap_co_values(template.__code__.co_names, mapping),
             co_varnames=_remap_co_values(template.__code__.co_varnames, mapping),
         ),
-        template.__globals__ | {f"__{field._name}_default__": field.type.__default__() for field in fields},
+        template.__globals__ | {f"__{field._name}_default__": _default_factory(field.type) for field in fields},
         argdefs=template.__defaults__,
     )
 
 
+def _default_factory(type_: type[BaseType]) -> Callable[[], Any]:
+    """Return a function that returns default values of a type, which are only shared if they're immutable."""
+    value = type_.__default__()
+    return repeat(value).__next__ if isinstance(value, _IMMUTABLE_TYPES) else type_.__default__
+
+
 def _generate_co_mapping(fields: list[Field]) -> dict[str, str]:
     """Generates a mapping of generated code object names to field names.
 

dissect/cstruct/types/wchar.py

@@ -28,12 +28,15 @@ class WcharArray(str, BaseArray):
         if cls.__null_terminated__:
             data += "\x00"
 
-        if not cls.__dynamic__ and (remaining := cls.__num_entries__ - (actual_size := len(data))):
+        buf = data.encode(Wchar.__encoding_map__[endian], "surrogatepass")
+
+        # A surrogate pair is a single character that takes up two entries of the array
+        if not cls.__dynamic__ and (remaining := cls.__num_entries__ - (actual_size := len(buf) // 2)):
             if remaining < 0:
                 raise ArraySizeError(f"Expected static array size {cls.__num_entries__}, got {actual_size} instead.")
-            data += "\x00" * remaining
+            buf += b"\x00\x00" * remaining
 
-        return stream.write(data.encode(Wchar.__encoding_map__[endian]))
+        return stream.write(buf)
 
 
 class Wchar(str, BaseType):
@@ -42,6 +45,7 @@ class Wchar(str, BaseType):
     __ArrayType__ = WcharArray
 
     __slots__ = ()
+    # Decode with "surrogatepass" so lone surrogates are handled correctly
     __encoding_map__: ClassVar[dict[Endianness, str]] = {
         "@": f"utf-16-{sys.byteorder[0]}e",
         "=": f"utf-16-{sys.byteorder[0]}e",
@@ -72,12 +76,13 @@ class Wchar(str, BaseType):
         if count != EOF and len(data) != count:
             raise EOFError(f"Read {len(data)} bytes, but expected {count}")
 
-        return type.__call__(cls, str(data, cls.__encoding_map__[endian]))
+        return type.__call__(cls, str(data, cls.__encoding_map__[endian], "surrogatepass"))
 
     @classmethod
     def _read_0(cls, stream: BinaryIO, *, context: dict[str, Any] | None = None, endian: Endianness) -> Wchar:
         if _is_seekable(stream):
-            return type.__call__(cls, _read_null_terminated(stream, 2).decode(cls.__encoding_map__[endian]))
+            data = _read_null_terminated(stream, 2)
+            return type.__call__(cls, data.decode(cls.__encoding_map__[endian], "surrogatepass"))
 
         buf = []
         while True:
@@ -90,8 +95,14 @@ class Wchar(str, BaseType):
 
             buf.append(point)
 
-        return type.__call__(cls, b"".join(buf).decode(cls.__encoding_map__[endian]))
+        return type.__call__(cls, b"".join(buf).decode(cls.__encoding_map__[endian], "surrogatepass"))
 
     @classmethod
     def _write(cls, stream: BinaryIO, data: str, *, endian: Endianness) -> int:
-        return stream.write(data.encode(cls.__encoding_map__[endian]))
+        buf = data.encode(cls.__encoding_map__[endian], "surrogatepass")
+
+        # A surrogate pair is a single character that doesn't fit in a single wchar
+        if (actual_size := len(buf) // 2) > 1:
+            raise ValueError(f"Expected a single character, got {actual_size} instead.")
+
+        return stream.write(buf or b"\x00\x00")

{dissect_cstruct-5.0.dev12.dist-info → dissect_cstruct-5.0.dev13.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dissect.cstruct
-Version: 5.0.dev12
+Version: 5.0.dev13
 Summary: A Dissect module implementing a parser for C-like structures: structure parsing in Python made easy
 Author-email: Dissect Team <dissect@fox-it.com>
 License-Expression: Apache-2.0

{dissect_cstruct-5.0.dev12.dist-info → dissect_cstruct-5.0.dev13.dist-info}/RECORD

@@ -1,7 +1,7 @@
 dissect/cstruct/__init__.py,sha256=KOyNeUvu4c0JaAiyOfj7iftGwUIjTmD8MwHtFsLPLSU,1504
 dissect/cstruct/bitbuffer.py,sha256=aNzY4lkWe8ffAc8j7FODXzwZc8gZwbey77UosyoYCsw,2711
-dissect/cstruct/compiler.py,sha256=hhX1bxZFsPOcs6BUO82YHNtWqhKsiOb-H-d18Aa1Q6E,35148
-dissect/cstruct/cstruct.py,sha256=VQDgGF7OUaPKQ_Wsy_2fUHL9Jlims6R2JOh-1_bgiZ8,26550
+dissect/cstruct/compiler.py,sha256=4oiJto_RqUvSOArScWo3XfgmiTkswxBgpu-qfQ9nYPA,36673
+dissect/cstruct/cstruct.py,sha256=mJerzPouWV29RdK5LqjGitDDMvNKcadoT2hj1zpExnM,26551
 dissect/cstruct/exception.py,sha256=RqD08LuBkbMqa3iQ1A05DNN8XgEmQR_SUnguwe9LGo8,316
 dissect/cstruct/expression.py,sha256=Y27-N48BomBoRFozbWhGDF_DS2lfklRwWVUbCzGnxHE,10873
 dissect/cstruct/lexer.py,sha256=NDUfAWgyxgRNRuVO8ExXQ51LqQXS2mv9bywVa_M7tFQ,19665
@@ -10,21 +10,21 @@ dissect/cstruct/util.py,sha256=vAf-kdOkk3zprW_gD6QmQeN3VIRgQqhPS4B_ha1jEpQ,18040
 dissect/cstruct/tools/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 dissect/cstruct/tools/stubgen.py,sha256=MeZkLmSeNSLNcAEwlAdIGv1lKvlvK3UQ4WN85ovsxTM,7672
 dissect/cstruct/types/__init__.py,sha256=6P71Pu94Dv7RNjZpDZ19pqzx-G0QzmQw4q0-3zWALso,911
-dissect/cstruct/types/base.py,sha256=10HLF2SDeWUMVc6ISMv8xKdEggqrxLPHnJMwI03DEFg,17216
-dissect/cstruct/types/char.py,sha256=Rzjn2-pXbn-fV_UnNRGgQugFZ50mPuXQhcmgHvY5PR0,3172
+dissect/cstruct/types/base.py,sha256=KOVgzQBsotj7xZoQcWnMzj5q56S3kUSnnNfXrtfTkqQ,17850
+dissect/cstruct/types/char.py,sha256=N19R7Zzn2uMuG-ULW5X2UgmgkLvjaBj1Td7AQQI93Fg,3417
 dissect/cstruct/types/enum.py,sha256=nOdI6vGtqTfPMhOBkKuiuOriIB2YHSt8Q9ruri37s3U,9177
 dissect/cstruct/types/flag.py,sha256=YxQHix8ZhK1qQX3J_Og7GJp_4JAxCwzUbUibuJeDaWQ,3479
 dissect/cstruct/types/int.py,sha256=9RbdrkvAI28zATM97VVORQnras0XhLaVMtdPjJwpTH4,2089
 dissect/cstruct/types/leb128.py,sha256=V97XTJT_Xbw1LpTwTN1OkjPHuQQAQUxd8OGuA6yEDmE,2387
 dissect/cstruct/types/packed.py,sha256=NgCX0YWX8SH8z47atq0f39w9MBjU_eBFJtUisUWdO-M,3867
 dissect/cstruct/types/pointer.py,sha256=FNuB9bfENWRebbKX-R6Vs6xWPMHepjvOhhomeUhr7I8,3985
-dissect/cstruct/types/structure.py,sha256=WDMmfhP2XE5p1Bn-WgFFstfaWC709AU7_y-SGi3TyDw,39116
+dissect/cstruct/types/structure.py,sha256=mxDvurQ0oAjI2flykhAPlT-03280cDpwGTz0-bqbr-0,39558
 dissect/cstruct/types/void.py,sha256=fV4hKk50uop_EKbR1nfy1cdtAzvUYrSymXsdakCyFk8,1222
-dissect/cstruct/types/wchar.py,sha256=qUbX6auj0qHWSOkgvaI7ZW_f7RGgwa-Hf_d_mDpntlM,3410
-dissect_cstruct-5.0.dev12.dist-info/licenses/COPYRIGHT,sha256=H-18RXfshdH9AdHwW2eO1Xa-5s6tY5eipHh5c0whDu4,316
-dissect_cstruct-5.0.dev12.dist-info/licenses/LICENSE,sha256=PhUqiw6jAh2KbBdVRPBq_hfAvfcTBin7nZ3CK7NQbTM,11341
-dissect_cstruct-5.0.dev12.dist-info/METADATA,sha256=TvTeDtkVFGsyMPWBgh3ANr5eKM1bLSi6SgI1hv4TUqw,9256
-dissect_cstruct-5.0.dev12.dist-info/WHEEL,sha256=mOdXR9C3wMVov7cSY8PY-XCzo3SmmWCHvGQr5csJoao,91
-dissect_cstruct-5.0.dev12.dist-info/entry_points.txt,sha256=z53zqZqwD2OLrAkRwrP4wTeiU9CQe7xrMly0T2c0_wQ,71
-dissect_cstruct-5.0.dev12.dist-info/top_level.txt,sha256=Mn-CQzEYsAbkxrUI0TnplHuXnGVKzxpDw_po_sXpvv4,8
-dissect_cstruct-5.0.dev12.dist-info/RECORD,,
+dissect/cstruct/types/wchar.py,sha256=VrnPXsP_lsyA-1xr4D2aSIVijumMhUuZWNlDHcn12-s,3964
+dissect_cstruct-5.0.dev13.dist-info/licenses/COPYRIGHT,sha256=H-18RXfshdH9AdHwW2eO1Xa-5s6tY5eipHh5c0whDu4,316
+dissect_cstruct-5.0.dev13.dist-info/licenses/LICENSE,sha256=PhUqiw6jAh2KbBdVRPBq_hfAvfcTBin7nZ3CK7NQbTM,11341
+dissect_cstruct-5.0.dev13.dist-info/METADATA,sha256=4sVLwtg9EZ3n5dRjzL-PqDvA1SQCr-JzbfoyTDcpcpM,9256
+dissect_cstruct-5.0.dev13.dist-info/WHEEL,sha256=mOdXR9C3wMVov7cSY8PY-XCzo3SmmWCHvGQr5csJoao,91
+dissect_cstruct-5.0.dev13.dist-info/entry_points.txt,sha256=z53zqZqwD2OLrAkRwrP4wTeiU9CQe7xrMly0T2c0_wQ,71
+dissect_cstruct-5.0.dev13.dist-info/top_level.txt,sha256=Mn-CQzEYsAbkxrUI0TnplHuXnGVKzxpDw_po_sXpvv4,8
+dissect_cstruct-5.0.dev13.dist-info/RECORD,,