Fix tests

Author: Schamper

dissect/cstruct/cstruct.py

@@ -220,7 +220,7 @@ def _next_anonymous(self) -> str:
         return name
 
     def _add_attr(self, name: str, value: Any, replace: bool = False) -> None:
-        if not replace and (name in self.__dict__ and self.__dict__[name] != value):
+        if not replace and ((existing := self.__dict__.get(name)) is not None and existing != value):
             raise ValueError(f"Attribute already exists: {name}")
         setattr(self, name, value)
 

dissect/cstruct/tools/stubgen.py

@@ -91,8 +91,8 @@ def generate_cstruct_stub(cs: cstruct, module_prefix: str = "", cls_name: str =
             stub = f"{name}: TypeAlias = {type_.__name__}"
         elif issubclass(type_, (types.Enum, types.Flag)):
             stub = generate_enum_stub(type_, cs_prefix=cs_prefix, module_prefix=module_prefix)
-        elif issubclass(typedef, types.Pointer):
-            typehint = generate_typehint(typedef, prefix=cs_prefix, module_prefix=module_prefix)
+        elif issubclass(type_, types.Pointer):
+            typehint = generate_typehint(type_, prefix=cs_prefix, module_prefix=module_prefix)
             stub = f"{name}: TypeAlias = {typehint}"
         elif issubclass(type_, types.Structure):
             stub = generate_structure_stub(type_, cs_prefix=cs_prefix, module_prefix=module_prefix)
@@ -167,9 +167,10 @@ def generate_structure_stub(
     module_prefix: str = "",
 ) -> str:
     result = [f"class {name_prefix}{structure.__name__}({module_prefix}{structure.__base__.__name__}):"]
-
     indent = " " * 4
 
+    all_types = structure.cs.typedefs | structure.cs.types
+
     args = ["self"]
     for field_name, field in structure.fields.items():
         type_name = field.type.__name__
@@ -180,7 +181,7 @@ def generate_structure_stub(
         while issubclass(nested_type, types.BaseArray):
             nested_type = nested_type.type
 
-        if issubclass(nested_type, types.Structure) and type_name not in structure.cs.typedefs:
+        if issubclass(nested_type, types.Structure) and type_name not in all_types:
             inlined = True
             inline_stub = generate_structure_stub(nested_type, cs_prefix=cs_prefix, module_prefix=module_prefix)
 

tests/test_annotations.py

@@ -8,12 +8,15 @@
 from dissect.cstruct.cstruct import cstruct
 from dissect.cstruct.types.base import BaseType
 
+CS = cstruct()
+
 
 @pytest.mark.parametrize(
     "name",
-    [name for name in cstruct().typedefs if " " not in name],
+    [name for name in CS.types | CS.typedefs if " " not in name],
 )
 def test_cstruct_type_annotation(name: str, monkeypatch: pytest.MonkeyPatch) -> None:
+    """Verify that all default types defined in cstruct have type annotations."""
     with (
         patch("typing.TYPE_CHECKING", True),
         patch("dissect.cstruct.types.base.MetaType.__getitem__", lambda self, item: self),

tests/test_basic.py

@@ -49,20 +49,20 @@ def test_load_init() -> None:
     """
     # load with first positional argument
     cs = cstruct(cdef)
-    assert "test" in cs.typedefs
+    assert "test" in cs.types
     assert cs.endian == "<"
 
     # load from keyword argument and big endian
     cs = cstruct(load=cdef, endian=">")
-    assert "test" in cs.typedefs
+    assert "test" in cs.types
     a = cs.test(a=0xBADC0DE, b=0xACCE55ED)
     assert len(bytes(a)) == 12
     assert bytes(a) == a.dumps()
     assert bytes(a) == b"\x0b\xad\xc0\xde\x00\x00\x00\x00\xac\xce\x55\xed"
 
     # load using positional argument and little endian
     cs = cstruct(cdef, endian="<")
-    assert "test" in cs.typedefs
+    assert "test" in cs.types
     a = cs.test(a=0xBADC0DE, b=0xACCE55ED)
     assert len(bytes(a)) == 12
     assert bytes(a) == a.dumps()
@@ -81,7 +81,7 @@ def test_load_init_kwargs_only() -> None:
         cs = cstruct(cdef, ">")
 
     cs = cstruct(cdef, endian=">")
-    assert "test" in cs.typedefs
+    assert "test" in cs.types
     assert cs.endian == ">"
 
 

tests/test_parser.py

@@ -188,7 +188,7 @@ def test_conditional_ifdef(cs: cstruct) -> None:
     """
     cs.load(cdef)
 
-    assert "test" in cs.typedefs
+    assert "test" in cs.types
 
 
 def test_conditional_ifndef(cs: cstruct) -> None:
@@ -218,7 +218,7 @@ def test_conditional_ifndef_guard(cs: cstruct) -> None:
     cs.load(cdef)
 
     assert "__MYGUARD" in cs.consts
-    assert "myStruct" in cs.typedefs
+    assert "myStruct" in cs.types
 
 
 def test_conditional_nested() -> None:
@@ -265,7 +265,7 @@ def test_conditional_in_struct(cs: cstruct) -> None:
     """
     cs.load(cdef)
 
-    assert "t_bitfield" in cs.typedefs
+    assert "t_bitfield" in cs.types
     assert "fval" in cs.t_bitfield.fields
     assert "bit0" in cs.t_bitfield.fields["fval"].type.fields
     assert "bit1" in cs.t_bitfield.fields["fval"].type.fields

tests/test_tools_stubgen.py

@@ -339,6 +339,35 @@ def __init__(self, fh: bytes | memoryview | bytearray | BinaryIO, /): ...
             """,
             id="pointer alias",
         ),
+        pytest.param(
+            """
+            struct A {
+                uint8 a;
+            };
+
+            struct Test {
+                A x;
+            };
+            """,
+            """
+            class cstruct(cstruct):
+                class A(Structure):
+                    a: cstruct.uint8
+                    @overload
+                    def __init__(self, a: cstruct.uint8 | None = ...): ...
+                    @overload
+                    def __init__(self, fh: bytes | memoryview | bytearray | BinaryIO, /): ...
+
+                class Test(Structure):
+                    x: cstruct.A
+                    @overload
+                    def __init__(self, x: cstruct.A | None = ...): ...
+                    @overload
+                    def __init__(self, fh: bytes | memoryview | bytearray | BinaryIO, /): ...
+
+            """,
+            id="reference other struct",
+        ),
     ],
 )
 def test_generate_cstruct_stub(cs: cstruct, cdef: str, expected: str) -> None: