Fix typedefs for internal types and add literals to cstruct-stubgen (#118)

Author: Schamper

dissect/cstruct/cstruct.py

@@ -425,10 +425,7 @@ def _make_union(
     ) -> type[Structure]:
         return self._make_struct(name, fields, align=align, anonymous=anonymous, base=Union)
 
-    Z = TYPE_CHECKING
-
     if TYPE_CHECKING:
-        A = 1
         # ruff: noqa: PYI042
         _int = int
         _float = float

dissect/cstruct/tools/stubgen.py

@@ -39,26 +39,29 @@ def generate_file_stub(path: Path, base: Path) -> str:
 
     header = [
         "# Generated by cstruct-stubgen",
-        "from typing import BinaryIO, overload",
+        "from typing import BinaryIO, Literal, overload",
         "",
         "import dissect.cstruct as __cs__",
+        "from typing_extensions import TypeAlias",
     ]
     body = []
 
     for name, obj in tmp_module.__dict__.items():
         if isinstance(obj, cstruct):
             stub = generate_cstruct_stub(obj, module_prefix="__cs__.", cls_name=f"_{name}")
             body.append(stub)
-            body.append(f"{name}: _{name}")
+
+            if body[-1][-1] != "\n":
+                body.append("")
+
+            body.append(f"# Technically `{name}` is an instance of `_{name}`, but then we can't use it in type hints")
+            body.append(f"{name}: TypeAlias = _{name}")
+            body.append("")
 
     if not body:
         return ""
 
-    body_str = "\n".join(body)
-    if "TypeAlias" in body_str:
-        header.append("from typing_extensions import TypeAlias")
-
-    return "\n".join([*header, "", body_str, ""])
+    return "\n".join([*header, "", "\n".join(body)])
 
 
 def generate_cstruct_stub(cs: cstruct, module_prefix: str = "", cls_name: str = "cstruct") -> str:
@@ -73,7 +76,7 @@ def generate_cstruct_stub(cs: cstruct, module_prefix: str = "", cls_name: str =
     for name, value in cs.consts.items():
         if name in empty_cs.consts:
             continue
-        body.append(textwrap.indent(f"{name}: {type(value).__name__} = ...", prefix=indent))
+        body.append(textwrap.indent(f"{name}: Literal[{value!r}] = ...", prefix=indent))
 
     defined_names = set()
 
@@ -82,10 +85,11 @@ def generate_cstruct_stub(cs: cstruct, module_prefix: str = "", cls_name: str =
         if name in empty_cs.typedefs:
             continue
 
-        if typedef.__name__ in defined_names:
+        if typedef.__name__ in empty_cs.typedefs:
+            stub = f"{name}: TypeAlias = {cs_prefix}{typedef.__name__}"
+        elif typedef.__name__ in defined_names:
             # Create an alias to the type if we have already seen it before.
             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.Structure):
@@ -129,7 +133,7 @@ def generate_generic_stub(
     cs_prefix: str = "",
     module_prefix: str = "",
 ) -> str:
-    return f"class {name_prefix}{type_.__name__}({module_prefix}{type_.__base__.__name__}): ..."
+    return f"class {name_prefix}{type_.__name__}({module_prefix}{type_.__base__.__name__}): ...\n"
 
 
 def generate_enum_stub(
@@ -140,6 +144,7 @@ def generate_enum_stub(
 ) -> str:
     result = [f"class {name_prefix}{enum.__name__}({module_prefix}{enum.__base__.__name__}):"]
     result.extend(f"    {key} = ..." for key in enum.__members__)
+    result.append("")
 
     return "\n".join(result)
 
@@ -181,6 +186,7 @@ def generate_structure_stub(
     result.append(
         textwrap.indent("def __init__(self, fh: bytes | memoryview | bytearray | BinaryIO, /): ...", prefix=indent)
     )
+    result.append("")
     return "\n".join(result)
 
 

tests/test_tools_stubgen.py

@@ -27,6 +27,7 @@
             class Test(Enum):
                 A = ...
                 B = ...
+
             """,
             id="enum",
         ),
@@ -41,6 +42,7 @@ class Test(Enum):
             class Test(Enum):
                 A = ...
                 B = ...
+
             """,
             id="enum int8",
         ),
@@ -55,6 +57,7 @@ class Test(Enum):
             class Test(Flag):
                 A = ...
                 B = ...
+
             """,
             id="flag",
         ),
@@ -63,7 +66,9 @@ class Test(Flag):
 def test_generate_enum_stub(cs: cstruct, cdef: str, expected: str) -> None:
     cs.load(cdef)
 
-    assert stubgen.generate_enum_stub(cs.Test) == textwrap.dedent(expected).strip()
+    # We don't want to strip all trailing whitespace in case it's part of the intended expected output
+    # So just remove one newline from the final """ block
+    assert stubgen.generate_enum_stub(cs.Test) == textwrap.dedent(expected).lstrip()[:-1]
 
 
 @pytest.mark.parametrize(
@@ -86,6 +91,7 @@ class Test(Structure):
                 def __init__(self, a: uint8 | None = ..., b: uint8 | None = ..., c: uint16 | None = ...): ...
                 @overload
                 def __init__(self, fh: bytes | memoryview | bytearray | BinaryIO, /): ...
+
             """,
             id="basic",
         ),
@@ -106,6 +112,7 @@ class Test(Structure):
                 def __init__(self, a: Array[uint8] | None = ..., b: CharArray | None = ..., c: WcharArray | None = ...): ...
                 @overload
                 def __init__(self, fh: bytes | memoryview | bytearray | BinaryIO, /): ...
+
             """,  # noqa: E501
             id="array",
         ),
@@ -124,6 +131,7 @@ class Test(Structure):
                 def __init__(self, a: Pointer[uint8] | None = ..., b: Array[Pointer[uint8]] | None = ...): ...
                 @overload
                 def __init__(self, fh: bytes | memoryview | bytearray | BinaryIO, /): ...
+
             """,
             id="pointer",
         ),
@@ -144,6 +152,7 @@ class Test(Structure):
                 def __init__(self, a: uint8 | None = ..., b: uint8 | None = ...): ...
                 @overload
                 def __init__(self, fh: bytes | memoryview | bytearray | BinaryIO, /): ...
+
             """,
             id="anonymous nested",
         ),
@@ -165,11 +174,13 @@ class __anonymous_0__(Union):
                     def __init__(self, a: uint8 | None = ..., b: uint8 | None = ...): ...
                     @overload
                     def __init__(self, fh: bytes | memoryview | bytearray | BinaryIO, /): ...
+
                 x: __anonymous_0__
                 @overload
                 def __init__(self, x: __anonymous_0__ | None = ...): ...
                 @overload
                 def __init__(self, fh: bytes | memoryview | bytearray | BinaryIO, /): ...
+
             """,
             id="named nested",
         ),
@@ -191,11 +202,13 @@ class __anonymous_0__(Structure):
                     def __init__(self, a: uint8 | None = ..., b: uint8 | None = ...): ...
                     @overload
                     def __init__(self, fh: bytes | memoryview | bytearray | BinaryIO, /): ...
+
                 x: Array[__anonymous_0__]
                 @overload
                 def __init__(self, x: Array[__anonymous_0__] | None = ...): ...
                 @overload
                 def __init__(self, fh: bytes | memoryview | bytearray | BinaryIO, /): ...
+
             """,
             id="named nested array",
         ),
@@ -204,7 +217,9 @@ def __init__(self, fh: bytes | memoryview | bytearray | BinaryIO, /): ...
 def test_generate_structure_stub(cs: cstruct, cdef: str, expected: str) -> None:
     cs.load(cdef)
 
-    assert stubgen.generate_structure_stub(cs.Test) == textwrap.dedent(expected).strip()
+    # We don't want to strip all trailing whitespace in case it's part of the intended expected output
+    # So just remove one newline from the final """ block
+    assert stubgen.generate_structure_stub(cs.Test) == textwrap.dedent(expected).lstrip()[:-1]
 
 
 @pytest.mark.parametrize(
@@ -225,10 +240,11 @@ def test_generate_structure_stub(cs: cstruct, cdef: str, expected: str) -> None:
             """,
             """
             class cstruct(cstruct):
-                TEST: int = ...
+                TEST: Literal[1] = ...
                 class TestEnum(Enum):
                     A = ...
                     B = ...
+
                 class TestStruct(Structure):
                     a: cstruct.uint8
                     @overload
@@ -253,16 +269,64 @@ class Test(Structure):
                     def __init__(self, a: cstruct.uint8 | None = ...): ...
                     @overload
                     def __init__(self, fh: bytes | memoryview | bytearray | BinaryIO, /): ...
+
                 _test: TypeAlias = Test
             """,
             id="alias stub",
         ),
+        pytest.param(
+            """
+            typedef __u16 __fs16;
+            typedef __u32 __fs32;
+            typedef __u64 __fs64;
+
+            struct Test {
+                __fs16 a;
+                __fs32 b;
+                __fs64 c;
+            };
+            """,
+            """
+            class cstruct(cstruct):
+                __fs16: TypeAlias = cstruct.uint16
+                __fs32: TypeAlias = cstruct.uint32
+                __fs64: TypeAlias = cstruct.uint64
+                class Test(Structure):
+                    a: cstruct.uint16
+                    b: cstruct.uint32
+                    c: cstruct.uint64
+                    @overload
+                    def __init__(self, a: cstruct.uint16 | None = ..., b: cstruct.uint32 | None = ..., c: cstruct.uint64 | None = ...): ...
+                    @overload
+                    def __init__(self, fh: bytes | memoryview | bytearray | BinaryIO, /): ...
+
+            """,  # noqa: E501
+            id="typedef stub",
+        ),
+        pytest.param(
+            """
+            #define INT 1
+            #define FLOAT 2.0
+            #define STRING "hello"
+            #define BYTES b'c'
+            """,
+            """
+            class cstruct(cstruct):
+                INT: Literal[1] = ...
+                FLOAT: Literal[2.0] = ...
+                STRING: Literal['hello'] = ...
+                BYTES: Literal[b'c'] = ...
+            """,
+            id="define literals",
+        ),
     ],
 )
 def test_generate_cstruct_stub(cs: cstruct, cdef: str, expected: str) -> None:
     cs.load(cdef)
 
-    assert stubgen.generate_cstruct_stub(cs) == textwrap.dedent(expected).strip()
+    # We don't want to strip all trailing whitespace in case it's part of the intended expected output
+    # So just remove one newline from the final """ block
+    assert stubgen.generate_cstruct_stub(cs) == textwrap.dedent(expected).lstrip()[:-1]
 
 
 def test_generate_cstruct_stub_empty(cs: cstruct) -> None:
@@ -292,9 +356,10 @@ def test_generate_file_stub(tmp_path: Path, monkeypatch: pytest.MonkeyPatch, cap
 
     expected = """
         # Generated by cstruct-stubgen
-        from typing import BinaryIO, overload
+        from typing import BinaryIO, Literal, overload
 
         import dissect.cstruct as __cs__
+        from typing_extensions import TypeAlias
 
         class _c_structure(__cs__.cstruct):
             class Test(__cs__.Structure):
@@ -304,7 +369,9 @@ class Test(__cs__.Structure):
                 def __init__(self, a: _c_structure.uint32 | None = ..., b: _c_structure.uint32 | None = ...): ...
                 @overload
                 def __init__(self, fh: bytes | memoryview | bytearray | BinaryIO, /): ...
-        c_structure: _c_structure
+
+        # Technically `c_structure` is an instance of `_c_structure`, but then we can't use it in type hints
+        c_structure: TypeAlias = _c_structure
     """
 
     assert stubgen.generate_file_stub(test_file, tmp_path) == textwrap.dedent(expected).lstrip()