Cleanup and move typedefs at the start of the file

dissect/cstruct/cstruct.py

@@ -408,24 +408,41 @@ def _make_union(
     ) -> type[Structure]:
         return self._make_struct(name, fields, align=align, anonymous=anonymous, base=Union)
 
-    def to_stub(self, name: str = ""):
+    def to_stub(self, name: str = "", ignore_type_defs: list[str] | None = None):
+        ignore_type_defs = ignore_type_defs or []
+
         buffer = io.StringIO()
         indentation = ""
         if name:
             buffer.write(f"class {name}(cstruct):\n")
             indentation = " " * 4
 
+        prev_offset = buffer.tell()
         for const, value in self.consts.items():
             buffer.write(indent(f"{const}: {type(value).__name__}=...\n", prefix=indentation))
 
+        buffer.write(self.stubify_typedefs(ignore_type_defs, indentation))
+
+        if prev_offset == buffer.tell():
+            buffer.write(indent("...", prefix=indentation))
+
+        output_value = buffer.getvalue()
+        buffer.close()
+        return output_value
+
+    def stubify_typedefs(self, ignore_type_defs: list[str] = None, indentation: str = "") -> str:
+        ignore_type_defs = ignore_type_defs or []
+        buffer = io.StringIO()
         for name, type_def in self.typedefs.items():
+            if name in ignore_type_defs:
+                continue
             if isinstance(type_def, MetaType) and (text := type_def.to_stub(name)):
                 buffer.write(indent(text, prefix=indentation))
                 buffer.write("\n")
 
-        output_value = buffer.getvalue()
+        output = buffer.getvalue()
         buffer.close()
-        return output_value
+        return output
 
 
 def ctypes(structure: Structure) -> _ctypes.Structure:

dissect/cstruct/tools/stubify.py

@@ -5,14 +5,15 @@
 import logging
 from argparse import ArgumentParser
 from pathlib import Path
+from types import ModuleType
 
 import dissect.cstruct.types as types
 from dissect.cstruct import cstruct
 
 log = logging.getLogger(__name__)
 
 
-def load_module(path: Path, base_path: Path):
+def load_module(path: Path, base_path: Path) -> ModuleType:
     module = None
     try:
         relative_path = path.relative_to(base_path)
@@ -26,25 +27,36 @@ def load_module(path: Path, base_path: Path):
         log.debug("Error while trying to import module %s", path, exc_info=e)
 
 
-def stubify_file(path: Path, base_path: Path):
-    buffer = io.StringIO()
-
-    cstruct_types = ",".join(types.__all__)
-    buffer.write(f"from dissect.cstruct.types import {cstruct_types}\n")
-    prev_offset = buffer.tell()
-
+def stubify_file(path: Path, base_path: Path) -> str:
     tmp_module = load_module(path, base_path)
     if tmp_module is None:
         return ""
 
+    if not hasattr(tmp_module, "cstruct"):
+        return ""
+
+    buffer = io.StringIO()
+    all_types = types.__all__.copy()
+    all_types.sort()
+
+    cstruct_types = ", ".join(all_types)
+    buffer.write("from __future__ import annotations\n\n")
+    buffer.write("from dissect.cstruct import cstruct\n")
+    buffer.write(f"from dissect.cstruct.types import {cstruct_types}\n\n")
+
+    empty_cstruct = cstruct()
+
+    buffer.write(empty_cstruct.stubify_typedefs())
+    buffer.write("\n")
+
+    prev_offset = buffer.tell()
+
     for name, variable in tmp_module.__dict__.items():
         if name.startswith("__"):
             continue
 
         if isinstance(variable, cstruct):
-            if variable._module != tmp_module.__name__:
-                continue
-            buffer.write(variable.to_stub(name))
+            buffer.write(variable.to_stub(name, empty_cstruct.typedefs.keys()))
 
     output = buffer.getvalue()
     if buffer.tell() == prev_offset:
@@ -85,6 +97,7 @@ def main():
                 continue
 
             with file.with_suffix(".pyi").open("wt") as output_file:
+                log.info(f"Writing stub of file {file} to {output_file}")
                 output_file.write(stub)