dissect.cstruct 4.6.dev4__tar.gz → 4.6.dev6__tar.gz

{dissect_cstruct-4.6.dev4 → dissect_cstruct-4.6.dev6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dissect.cstruct
-Version: 4.6.dev4
+Version: 4.6.dev6
 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: Apache License 2.0

{dissect_cstruct-4.6.dev4 → dissect_cstruct-4.6.dev6}/dissect/cstruct/parser.py

@@ -314,11 +314,11 @@ class TokenParser(Parser):
                 tokens.eol()
                 break
 
-            if tokens.next not in (self.TOK.NAME, self.TOK.DEFS):
+            if tokens.next not in (self.TOK.NAME, self.TOK.DEFS, self.TOK.IDENTIFIER):
                 break
 
             ntoken = tokens.consume()
-            if ntoken == self.TOK.NAME:
+            if ntoken in (self.TOK.NAME, self.TOK.IDENTIFIER):
                 names.append(ntoken.value.strip())
             elif ntoken == self.TOK.DEFS:
                 names.extend([name.strip() for name in ntoken.value.strip().split(",")])

{dissect_cstruct-4.6.dev4 → dissect_cstruct-4.6.dev6}/dissect/cstruct/tools/stubgen.py

@@ -92,6 +92,9 @@ def generate_cstruct_stub(cs: cstruct, module_prefix: str = "", cls_name: str =
             stub = f"{name}: TypeAlias = {typedef.__name__}"
         elif issubclass(typedef, (types.Enum, types.Flag)):
             stub = generate_enum_stub(typedef, cs_prefix=cs_prefix, module_prefix=module_prefix)
+        elif issubclass(typedef, types.Pointer):
+            typehint = generate_typehint(typedef, prefix=cs_prefix, module_prefix=module_prefix)
+            stub = f"{name}: TypeAlias = {typehint}"
         elif issubclass(typedef, types.Structure):
             stub = generate_structure_stub(typedef, cs_prefix=cs_prefix, module_prefix=module_prefix)
         elif issubclass(typedef, types.BaseType):

{dissect_cstruct-4.6.dev4 → dissect_cstruct-4.6.dev6}/dissect/cstruct/types/structure.py

@@ -828,13 +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 = _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(chain.from_iterable(zip((f"__{name}_default__" for name in field_names), field_names))),
-            co_varnames=("self", *field_names),
+            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__,
@@ -847,20 +847,50 @@ def _generate_union__init__(fields: list[Field]) -> FunctionType:
     Args:
         fields: List of field names.
     """
-    field_names = [field._name for field in fields]
+    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=(None, *field_names),
-            co_names=("object", "__setattr__", *(f"__{name}_default__" for name in field_names)),
-            co_varnames=("self", *field_names),
+            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.
 

{dissect_cstruct-4.6.dev4 → dissect_cstruct-4.6.dev6}/dissect.cstruct.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dissect.cstruct
-Version: 4.6.dev4
+Version: 4.6.dev6
 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: Apache License 2.0

{dissect_cstruct-4.6.dev4 → dissect_cstruct-4.6.dev6}/tests/test_parser.py

@@ -7,7 +7,7 @@ import pytest
 
 from dissect.cstruct.exceptions import ParserError
 from dissect.cstruct.parser import TokenParser
-from dissect.cstruct.types import BaseArray, Pointer
+from dissect.cstruct.types import BaseArray, Pointer, Structure
 from tests.utils import verify_compiled
 
 if TYPE_CHECKING:
@@ -149,3 +149,18 @@ def test_includes(cs: cstruct) -> None:
     assert cs.myStruct.__name__ == "myStruct"
     assert len(cs.myStruct.fields) == 1
     assert cs.myStruct.fields.get("charVal")
+
+
+def test_typedef_pointer(cs: cstruct) -> None:
+    cdef = """
+    typedef struct _IMAGE_DATA_DIRECTORY {
+        DWORD VirtualAddress;
+        DWORD Size;
+    } IMAGE_DATA_DIRECTORY, *PIMAGE_DATA_DIRECTORY;
+    """
+    cs.load(cdef)
+
+    assert issubclass(cs._IMAGE_DATA_DIRECTORY, Structure)
+    assert cs.IMAGE_DATA_DIRECTORY is cs._IMAGE_DATA_DIRECTORY
+    assert issubclass(cs.PIMAGE_DATA_DIRECTORY, Pointer)
+    assert cs.PIMAGE_DATA_DIRECTORY.type == cs._IMAGE_DATA_DIRECTORY

{dissect_cstruct-4.6.dev4 → dissect_cstruct-4.6.dev6}/tests/test_tools_stubgen.py

@@ -319,6 +319,26 @@ def test_generate_structure_stub(cs: cstruct, cdef: str, expected: str) -> None:
             """,
             id="define literals",
         ),
+        pytest.param(
+            """
+            typedef struct _Test {
+                uint8 a;
+            } Test, *pTest;
+            """,
+            """
+            class cstruct(cstruct):
+                class _Test(Structure):
+                    a: cstruct.uint8
+                    @overload
+                    def __init__(self, a: cstruct.uint8 | None = ...): ...
+                    @overload
+                    def __init__(self, fh: bytes | memoryview | bytearray | BinaryIO, /): ...
+
+                Test: TypeAlias = _Test
+                pTest: TypeAlias = Pointer[cstruct._Test]
+            """,
+            id="pointer alias",
+        ),
     ],
 )
 def test_generate_cstruct_stub(cs: cstruct, cdef: str, expected: str) -> None: