Create a tool for stubs in cstruct v4

dissect/cstruct/tools/stubify.py

@@ -0,0 +1,184 @@
+from __future__ import annotations
+
+import importlib
+import importlib.util
+import logging
+from argparse import ArgumentParser
+from pathlib import Path
+from textwrap import indent
+from types import FunctionType, ModuleType
+from typing import TYPE_CHECKING, Any
+
+import dissect.cstruct.types as types
+from dissect.cstruct import cstruct
+
+if TYPE_CHECKING:
+    from collections.abc import Iterable
+
+
+log = logging.getLogger(__name__)
+
+
+def load_module(path: Path, base_path: Path) -> ModuleType | None:
+    module = None
+    try:
+        relative_path = path.relative_to(base_path)
+        module_tuple = (*relative_path.parent.parts, relative_path.stem)
+        spec = importlib.util.spec_from_file_location(".".join(module_tuple), path)
+        module = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(module)
+    except Exception as e:
+        log.exception("Unable to import %s", path)
+        log.debug("Error while trying to import module %s", path, exc_info=e)
+
+    return module
+
+
+def to_type(_type: type | Any) -> type:
+    if not isinstance(_type, type):
+        _type = type(_type)
+    return _type
+
+
+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 ""
+
+    all_types = types.__all__.copy()
+    all_types.sort()
+    all_types.append("")
+
+    cstruct_types = indent(",\n".join(all_types), prefix=" " * 4)
+    result = [
+        "from __future__ import annotations\n",
+        "from typing import overload, BinaryIO\n",
+        "from typing_extensions import TypeAlias\n",
+        "from dissect.cstruct import cstruct",
+        f"from dissect.cstruct.types import (\n{cstruct_types})\n",
+    ]
+
+    prev_entries = len(result)
+
+    for name, variable in tmp_module.__dict__.items():
+        if name.startswith("__"):
+            continue
+
+        if isinstance(variable, ModuleType):
+            result.append(f"import {name}")
+        elif isinstance(variable, cstruct):
+            result.append(stubify_cstruct(variable, name))
+        elif isinstance(variable, (bytes, bytearray, str, int, float, dict, list, tuple)):
+            result.append(f"{name}: {type(variable).__name__}")
+        elif isinstance(variable, FunctionType):
+            anno = variable.__annotations__
+            _items = list(anno.items())[:-1]
+
+            args = ", ".join(f"{name}: {to_type(_type).__name__}" for (name, _type) in _items)
+
+            return_value = repr(to_type(anno.get("return")).__name__)
+            signature = f"def {name}({''.join(args)}) -> {return_value}:"
+            if variable.__doc__:
+                result.append(signature)
+                result.append(f'    """{variable.__doc__}"""')
+                result.append("    ...")
+            else:
+                result.append(f"{signature} ...")
+        elif "dissect.cstruct" in variable.__module__:
+            if hasattr(variable, "cs"):
+                result.append(f"{name}: {variable.cs.__type_def_name__}.{variable.__name__}")
+        elif isinstance(variable, type):
+            result.append(f"from {variable.__module__} import {name}")
+        else:
+            result.append(f"{name}: {variable.__class__.__name__}")
+
+    if prev_entries == len(result):
+        return ""
+
+    # Empty line at the end of the file
+    result.append("")
+    return "\n".join(result)
+
+
+def stubify_cstruct(c_structure: cstruct, name: str = "", ignore_type_defs: Iterable[str] | None = None) -> str:
+    ignore_type_defs = ignore_type_defs or []
+
+    result = []
+    indentation = ""
+    if name:
+        result.append(f"class {name}(cstruct):")
+        indentation = " " * 4
+        c_structure.__type_def_name__ = name
+
+    prev_length = len(result)
+    for const, value in c_structure.consts.items():
+        result.append(indent(f"{const}: {type(value).__name__} = ...", prefix=indentation))
+
+    if type_defs := stubify_typedefs(c_structure, ignore_type_defs, indentation):
+        result.append(type_defs)
+
+    if prev_length == len(result):
+        # an empty definition, add elipses
+        result.append(indent("...", prefix=indentation))
+
+    return "\n".join(result)
+
+
+def stubify_typedefs(c_structure: cstruct, ignore_type_defs: Iterable[str] | None = None, indentation: str = "") -> str:
+    ignore_type_defs = ignore_type_defs or []
+
+    result = []
+    for name, type_def in c_structure.typedefs.items():
+        if name in ignore_type_defs:
+            continue
+
+        if isinstance(type_def, types.MetaType) and (text := type_def.to_type_stub(name)):
+            result.append(indent(text, prefix=indentation))
+
+    return "\n".join(result)
+
+
+def setup_logger(verbosity: int) -> None:
+    level = logging.INFO
+    if verbosity >= 1:
+        level = logging.DEBUG
+
+    logging.basicConfig(level=level)
+
+
+def main() -> None:
+    description = """
+        Create stub files for cstruct definitions.
+
+        These stub files are in a `.pyi` format and provides type information to cstruct definitions.
+    """
+
+    parser = ArgumentParser("stubify", description=description)
+    parser.add_argument("path", type=Path)
+    parser.add_argument("-v", "--verbose", action="count", default=0)
+    args = parser.parse_args()
+
+    setup_logger(args.verbose)
+
+    file_path: Path = args.path
+
+    iterator = file_path.rglob("*.py")
+    if file_path.is_file():
+        iterator = [file_path]
+
+    for file in iterator:
+        if file.is_file() and file.suffix == ".py":
+            stub = stubify_file(file, file_path)
+            if not stub:
+                continue
+
+            with file.with_suffix(".pyi").open("wt") as output_file:
+                log.info("Writing stub of file %s to %s", file, output_file.name)
+                output_file.write(stub)
+
+
+if __name__ == "__main__":
+    main()

dissect/cstruct/types/base.py

@@ -178,6 +178,22 @@
         """
         return cls._write_array(stream, [*array, cls.__default__()])
 
+    def _class_stub(cls) -> str:
+        return f"class {cls.__name__}({cls.__base__.__name__}):"
+
+    def _type_stub(cls, name: str = "", underscore: bool = False) -> str:
+        cls_name = cls.__name__
+        if underscore:
+            cls_name = f"_{cls_name}"
+
+        if cls.__name__ in cls.cs.typedefs and (cs_name := getattr(cls.cs, "__type_def_name__", "")):
+            cls_name = f"{cs_name}.{cls_name}"
+
+        return f"{name}: {cls_name}"
+
+    def to_type_stub(cls, name: str) -> str:
+        return ""
+
 
 class _overload:
     """Descriptor to use on the ``write`` and ``dumps`` methods on cstruct types.
@@ -244,6 +260,14 @@
 
         return cls.type._read_array(stream, num, context)
 
+    def default(cls) -> BaseType:
+        return type.__call__(
+            cls, [cls.type.default() for _ in range(0 if cls.dynamic or cls.null_terminated else cls.num_entries)]
+        )
+
+    def _type_stub(cls, name: str = "", underscore: bool = False) -> str:
+        return f"{name}: {cls.__base__.__name__}"
+
 
 class Array(list, BaseType, metaclass=ArrayMetaType):
     """Implements a fixed or dynamically sized array type.
@@ -270,6 +294,15 @@
 
         return cls.type._write_array(stream, data)
 
+    @classmethod
+    def _type_stub(cls, name: str = "", underscore: bool = False) -> str:
+        cls_name = cls.type.__name__
+
+        if cls_name in cls.cs.typedefs and (cs_name := getattr(cls.cs, "__type_def_name__", "")):
+            cls_name = f"{cs_name}.{cls_name}"
+
+        return f"{name}: {cls.__base__.__name__}[{cls_name}]"
+
 
 def _is_readable_type(value: Any) -> bool:
     return hasattr(value, "read")

dissect/cstruct/types/enum.py

@@ -84,6 +84,15 @@ def _write_0(cls, stream: BinaryIO, array: list[BaseType]) -> int:
         data = [entry.value if isinstance(entry, Enum) else entry for entry in array]
         return cls._write_array(stream, [*data, cls.type.__default__()])
 
+    def _class_stub(cls) -> str:
+        return f"class {cls.__name__}({cls.__base__.__name__}, {cls.type.__name__}):"
+
+    def to_type_stub(cls, name: str = "") -> str:
+        result = [cls._class_stub()]
+        result.extend(f"    {key} = ..." for key in cls.__members__)
+
+        return "\n".join(result)
+
 
 def _fix_alias_members(cls: type[Enum]) -> None:
     # Emulate aenum NoAlias behaviour

dissect/cstruct/types/packed.py

@@ -2,7 +2,7 @@
 
 from functools import lru_cache
 from struct import Struct
-from typing import Any, BinaryIO
+from typing import Any, BinaryIO, Generic, TypeVar
 
 from dissect.cstruct.types.base import EOF, BaseType
 
@@ -12,17 +12,20 @@ def _struct(endian: str, packchar: str) -> Struct:
     return Struct(f"{endian}{packchar}")
 
 
-class Packed(BaseType):
+T = TypeVar("T", int, float)
+
+
+class Packed(BaseType, Generic[T]):
     """Packed type for Python struct (un)packing."""
 
     packchar: str
 
     @classmethod
-    def _read(cls, stream: BinaryIO, context: dict[str, Any] | None = None) -> Packed:
+    def _read(cls, stream: BinaryIO, context: dict[str, Any] | None = None) -> Packed[T]:
         return cls._read_array(stream, 1, context)[0]
 
     @classmethod
-    def _read_array(cls, stream: BinaryIO, count: int, context: dict[str, Any] | None = None) -> list[Packed]:
+    def _read_array(cls, stream: BinaryIO, count: int, context: dict[str, Any] | None = None) -> list[Packed[T]]:
         if count == EOF:
             data = stream.read()
             length = len(data)
@@ -39,7 +42,7 @@ def _read_array(cls, stream: BinaryIO, count: int, context: dict[str, Any] | Non
         return [cls.__new__(cls, value) for value in fmt.unpack(data)]
 
     @classmethod
-    def _read_0(cls, stream: BinaryIO, context: dict[str, Any] | None = None) -> Packed:
+    def _read_0(cls, stream: BinaryIO, context: dict[str, Any] | None = None) -> Packed[T]:
         result = []
 
         fmt = _struct(cls.cs.endian, cls.packchar)
@@ -57,9 +60,13 @@ def _read_0(cls, stream: BinaryIO, context: dict[str, Any] | None = None) -> Pac
         return result
 
     @classmethod
-    def _write(cls, stream: BinaryIO, data: Packed) -> int:
+    def _write(cls, stream: BinaryIO, data: Packed[T]) -> int:
         return stream.write(_struct(cls.cs.endian, cls.packchar).pack(data))
 
     @classmethod
-    def _write_array(cls, stream: BinaryIO, data: list[Packed]) -> int:
+    def _write_array(cls, stream: BinaryIO, data: list[Packed[T]]) -> int:
         return stream.write(_struct(cls.cs.endian, f"{len(data)}{cls.packchar}").pack(*data))
+
+    @classmethod
+    def to_type_stub(cls, name: str) -> str:
+        return f"{name}: TypeAlias = Packed[{cls.__base__.__name__}]"

dissect/cstruct/types/pointer.py

@@ -1,22 +1,24 @@
 from __future__ import annotations
 
-from typing import Any, BinaryIO
+from typing import Any, BinaryIO, Generic, TypeVar
 
 from dissect.cstruct.exceptions import NullPointerDereference
 from dissect.cstruct.types.base import BaseType, MetaType
 from dissect.cstruct.types.char import Char
 from dissect.cstruct.types.void import Void
 
+T = TypeVar("T", bound=MetaType)
 
-class Pointer(int, BaseType):
+
+class Pointer(int, BaseType, Generic[T]):
     """Pointer to some other type."""
 
-    type: MetaType
+    type: T
     _stream: BinaryIO | None
     _context: dict[str, Any] | None
     _value: BaseType
 
-    def __new__(cls, value: int, stream: BinaryIO | None, context: dict[str, Any] | None = None) -> Pointer:  # noqa: PYI034
+    def __new__(cls, value: int, stream: BinaryIO | None, context: dict[str, Any] | None = None) -> Pointer[T]:
         obj = super().__new__(cls, value)
         obj._stream = stream
         obj._context = context
@@ -70,15 +72,15 @@ def __default__(cls) -> Pointer:
         return cls.__new__(cls, cls.cs.pointer.__default__(), None, None)
 
     @classmethod
-    def _read(cls, stream: BinaryIO, context: dict[str, Any] | None = None) -> Pointer:
+    def _read(cls, stream: BinaryIO, context: dict[str, Any] | None = None) -> Pointer[T]:
         return cls.__new__(cls, cls.cs.pointer._read(stream, context), stream, context)
 
     @classmethod
     def _write(cls, stream: BinaryIO, data: int) -> int:
         return cls.cs.pointer._write(stream, data)
 
-    def dereference(self) -> Any:
-        if self == 0 or self._stream is None:
+    def dereference(self) -> T:
+        if self == 0:
             raise NullPointerDereference
 
         if self._value is None and not issubclass(self.type, Void):
@@ -97,3 +99,7 @@ def dereference(self) -> Any:
             self._value = value
 
         return self._value
+
+    @classmethod
+    def _type_stub(cls, name: str = "", underscore: bool = False) -> str:
+        return f"{name}: {cls.__base__.__name__}[{cls.type.__name__}]"