From 6277706d4baa3226eec39467a887284139877e9d Mon Sep 17 00:00:00 2001 From: MrYoranimo <1333535+MrYoranimo@users.noreply.github.com> Date: Mon, 28 Jul 2025 13:46:45 +0200 Subject: [PATCH 1/2] Refactor structure and union initializer generator --- dissect/cstruct/types/structure.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/dissect/cstruct/types/structure.py b/dissect/cstruct/types/structure.py index 1234d24..eca5f85 100644 --- a/dissect/cstruct/types/structure.py +++ b/dissect/cstruct/types/structure.py @@ -829,12 +829,14 @@ def _generate_structure__init__(fields: list[Field]) -> FunctionType: fields: List of field names. """ field_names = [field._name for field in fields] + mapping = {f"_{i}": name for i, name in enumerate(field_names)} + mapping.update({f"_{i}_default": f"__{name}_default__" for i, name in enumerate(field_names)}) template: FunctionType = _make_structure__init__(len(field_names)) return type(template)( template.__code__.replace( - co_names=tuple(chain.from_iterable(zip((f"__{name}_default__" for name in field_names), field_names))), - co_varnames=("self", *field_names), + co_names=tuple(mapping.get(a, a) for a in template.__code__.co_names), + co_varnames=tuple(mapping.get(v, v) for v in template.__code__.co_varnames), ), template.__globals__ | {f"__{field._name}_default__": field.type.__default__() for field in fields}, argdefs=template.__defaults__, @@ -848,13 +850,15 @@ def _generate_union__init__(fields: list[Field]) -> FunctionType: fields: List of field names. """ field_names = [field._name for field in fields] + field_mapping = {f"_{i}": name for i, name in enumerate(field_names)} + defaults_mapping = {f"_{i}_default": f"__{name}_default__" for i, name in enumerate(field_names)} template: FunctionType = _make_union__init__(len(field_names)) return type(template)( template.__code__.replace( - co_consts=(None, *field_names), - co_names=("object", "__setattr__", *(f"__{name}_default__" for name in field_names)), - co_varnames=("self", *field_names), + co_consts=tuple(field_mapping.get(c, c) for c in template.__code__.co_consts), + co_names=tuple(defaults_mapping.get(a, a) for a in template.__code__.co_names), + co_varnames=tuple(field_mapping.get(v, v) for v in template.__code__.co_varnames), ), template.__globals__ | {f"__{field._name}_default__": field.type.__default__() for field in fields}, argdefs=template.__defaults__, From aef481aac06d85320b45186bf0cf0b2a2d9cee9f Mon Sep 17 00:00:00 2001 From: Schamper <1254028+Schamper@users.noreply.github.com> Date: Mon, 28 Jul 2025 17:28:00 +0200 Subject: [PATCH 2/2] Slight changes --- dissect/cstruct/types/structure.py | 52 ++++++++++++++++++++++-------- 1 file changed, 39 insertions(+), 13 deletions(-) diff --git a/dissect/cstruct/types/structure.py b/dissect/cstruct/types/structure.py index eca5f85..8d6d61b 100644 --- a/dissect/cstruct/types/structure.py +++ b/dissect/cstruct/types/structure.py @@ -828,15 +828,13 @@ def _generate_structure__init__(fields: list[Field]) -> FunctionType: Args: fields: List of field names. """ - field_names = [field._name for field in fields] - mapping = {f"_{i}": name for i, name in enumerate(field_names)} - mapping.update({f"_{i}_default": f"__{name}_default__" for i, name in enumerate(field_names)}) + mapping = _generate_co_mapping(fields) - template: FunctionType = _make_structure__init__(len(field_names)) + template: FunctionType = _make_structure__init__(len(fields)) return type(template)( template.__code__.replace( - co_names=tuple(mapping.get(a, a) for a in template.__code__.co_names), - co_varnames=tuple(mapping.get(v, v) for v in template.__code__.co_varnames), + 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}, argdefs=template.__defaults__, @@ -849,22 +847,50 @@ def _generate_union__init__(fields: list[Field]) -> FunctionType: Args: fields: List of field names. """ - field_names = [field._name for field in fields] - field_mapping = {f"_{i}": name for i, name in enumerate(field_names)} - defaults_mapping = {f"_{i}_default": f"__{name}_default__" for i, name in enumerate(field_names)} + mapping = _generate_co_mapping(fields) - template: FunctionType = _make_union__init__(len(field_names)) + template: FunctionType = _make_union__init__(len(fields)) return type(template)( template.__code__.replace( - co_consts=tuple(field_mapping.get(c, c) for c in template.__code__.co_consts), - co_names=tuple(defaults_mapping.get(a, a) for a in template.__code__.co_names), - co_varnames=tuple(field_mapping.get(v, v) for v in template.__code__.co_varnames), + co_consts=_remap_co_values(template.__code__.co_consts, mapping), + 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}, argdefs=template.__defaults__, ) +def _generate_co_mapping(fields: list[Field]) -> dict[str, str]: + """Generates a mapping of generated code object names to field names. + + The generated code uses names like ``_0``, ``_1``, etc. for fields, and ``_0_default``, ``_1_default``, etc. + for default initializer values. Return a mapping of these names to the actual field names. + + Args: + fields: List of field names. + """ + return { + key: value + for i, field in enumerate(fields) + for key, value in [(f"_{i}", field._name), (f"_{i}_default", f"__{field._name}_default__")] + } + + +def _remap_co_values(value: tuple[Any, ...], mapping: dict[str, str]) -> tuple[Any, ...]: + """Remaps code object values using a mapping. + + This is used to replace generated code object names with actual field names. + + Args: + value: The original code object values. + mapping: A mapping of generated code object names to field names. + """ + # Only attempt to remap if the value is a string, otherwise return it as is + # This is to avoid issues with trying to remap non-hashable types, and we only need to replace strings anyway + return tuple(mapping.get(v, v) if isinstance(v, str) else v for v in value) + + def _generate__eq__(fields: list[str]) -> FunctionType: """Generates an ``__eq__`` method for a class with the specified fields.