dissect.cstruct 4.4.dev3__py3-none-any.whl → 4.5.dev2__py3-none-any.whl

dissect/cstruct/tools/stubgen.py

@@ -0,0 +1,226 @@
+from __future__ import annotations
+
+import argparse
+import importlib
+import importlib.util
+import logging
+import textwrap
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from dissect.cstruct import types
+from dissect.cstruct.cstruct import cstruct
+
+if TYPE_CHECKING:
+    from types import ModuleType
+
+log = logging.getLogger(__name__)
+
+
+def load_module(path: Path, base: Path) -> ModuleType | None:
+    module = None
+    try:
+        relative_path = path.relative_to(base)
+        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.warning("Unable to import %s", path)
+        log.debug("Error while trying to import module %s", path, exc_info=e)
+
+    return module
+
+
+def generate_file_stub(path: Path, base: Path) -> str:
+    tmp_module = load_module(path, base)
+    if tmp_module is None or not hasattr(tmp_module, "cstruct"):
+        return ""
+
+    header = [
+        "# Generated by cstruct-stubgen",
+        "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)
+
+            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 ""
+
+    return "\n".join([*header, "", "\n".join(body)])
+
+
+def generate_cstruct_stub(cs: cstruct, module_prefix: str = "", cls_name: str = "cstruct") -> str:
+    empty_cs = cstruct()
+
+    cs_prefix = f"{cls_name}."
+    header = [f"class {cls_name}({module_prefix}cstruct):"]
+    body = []
+    indent = " " * 4
+
+    # Constants first
+    for name, value in cs.consts.items():
+        if name in empty_cs.consts:
+            continue
+        body.append(textwrap.indent(f"{name}: Literal[{value!r}] = ...", prefix=indent))
+
+    defined_names = set()
+
+    # Then typedefs
+    for name, typedef in cs.typedefs.items():
+        if name in empty_cs.typedefs:
+            continue
+
+        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):
+            stub = generate_structure_stub(typedef, cs_prefix=cs_prefix, module_prefix=module_prefix)
+        elif issubclass(typedef, types.BaseType):
+            stub = generate_generic_stub(typedef, cs_prefix=cs_prefix, module_prefix=module_prefix)
+        elif isinstance(typedef, str):
+            stub = f"{name}: TypeAlias = {typedef}"
+        else:
+            raise TypeError(f"Unknown typedef: {typedef}")
+
+        defined_names.add(typedef.__name__)
+
+        body.append(textwrap.indent(stub, prefix=indent))
+
+    if not body:
+        body.append(textwrap.indent("...", prefix=indent))
+
+    return "\n".join(header + body)
+
+
+def generate_typehint(
+    type_: type[types.BaseType],
+    prefix: str = "",
+    module_prefix: str = "",
+) -> str:
+    if issubclass(type_, types.CharArray):
+        return f"{module_prefix}CharArray"
+    if issubclass(type_, types.WcharArray):
+        return f"{module_prefix}WcharArray"
+    if issubclass(type_, types.Pointer):
+        return f"{module_prefix}Pointer[{generate_typehint(type_.type, prefix, module_prefix)}]"
+    if issubclass(type_, types.Array):
+        return f"{module_prefix}Array[{generate_typehint(type_.type, prefix, module_prefix)}]"
+    return f"{prefix}{type_.__name__}"
+
+
+def generate_generic_stub(
+    type_: type[types.BaseType],
+    name_prefix: str = "",
+    cs_prefix: str = "",
+    module_prefix: str = "",
+) -> str:
+    return f"class {name_prefix}{type_.__name__}({module_prefix}{type_.__base__.__name__}): ...\n"
+
+
+def generate_enum_stub(
+    enum: type[types.Enum | types.Flag],
+    name_prefix: str = "",
+    cs_prefix: str = "",
+    module_prefix: str = "",
+) -> 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)
+
+
+def generate_structure_stub(
+    structure: type[types.Structure],
+    name_prefix: str = "",
+    cs_prefix: str = "",
+    module_prefix: str = "",
+) -> str:
+    result = [f"class {name_prefix}{structure.__name__}({module_prefix}{structure.__base__.__name__}):"]
+
+    indent = " " * 4
+
+    args = ["self"]
+    for field_name, field in structure.fields.items():
+        type_name = field.type.__name__
+        inlined = False
+
+        # If it's a structure and not globally defined, add an inline stub for it
+        nested_type = field.type
+        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:
+            inlined = True
+            inline_stub = generate_structure_stub(nested_type, cs_prefix=cs_prefix, module_prefix=module_prefix)
+
+            result.append(textwrap.indent(inline_stub, prefix=indent))
+
+        type_hint = generate_typehint(field.type, "" if inlined else f"{cs_prefix}", module_prefix)
+        result.append(f"    {field_name}: {type_hint}")
+
+        args.append(f"{field_name}: {type_hint} | None = ...")
+
+    result.append(textwrap.indent("@overload", prefix=indent))
+    result.append(textwrap.indent(f"def __init__({', '.join(args)}): ...", prefix=indent))
+    result.append(textwrap.indent("@overload", prefix=indent))
+    result.append(
+        textwrap.indent("def __init__(self, fh: bytes | memoryview | bytearray | BinaryIO, /): ...", prefix=indent)
+    )
+    result.append("")
+    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:
+    parser = argparse.ArgumentParser(
+        "cstruct-stubify",
+        description="Create .pyi stub files for cstruct definitions",
+        epilog="NOTE: This tool will only generate stubs for the cstruct definitions in a file, not any other Python code. Manual fixups may be required.",  # noqa: E501
+    )
+    parser.add_argument("path", type=Path, help="path to the file or directory to create stubs for")
+    parser.add_argument("-v", "--verbose", action="count", default=0)
+    args = parser.parse_args()
+
+    setup_logger(args.verbose)
+
+    path: Path = args.path
+    for file in path.rglob("*.py") if path.is_dir() else [path]:
+        if file.is_file() and file.suffix == ".py":
+            stub = generate_file_stub(file, path)
+            if not stub:
+                continue
+
+            stub_file = file.with_suffix(".pyi")
+            log.info("Writing stub of file %s to %s", file, stub_file.name)
+            stub_file.write_text(stub)
+
+
+if __name__ == "__main__":
+    main()

dissect/cstruct/types/__init__.py

@@ -1,4 +1,4 @@
-from dissect.cstruct.types.base import Array, ArrayMetaType, BaseType, MetaType
+from dissect.cstruct.types.base import Array, BaseArray, BaseType, MetaType
 from dissect.cstruct.types.char import Char, CharArray
 from dissect.cstruct.types.enum import Enum
 from dissect.cstruct.types.flag import Flag
@@ -13,7 +13,7 @@ from dissect.cstruct.types.wchar import Wchar, WcharArray
 __all__ = [
     "LEB128",
     "Array",
-    "ArrayMetaType",
+    "BaseArray",
     "BaseType",
     "Char",
     "CharArray",

dissect/cstruct/types/base.py

@@ -2,12 +2,14 @@ from __future__ import annotations
 
 import functools
 from io import BytesIO
-from typing import TYPE_CHECKING, Any, BinaryIO, Callable
+from typing import TYPE_CHECKING, Any, BinaryIO, Callable, ClassVar, TypeVar
 
 from dissect.cstruct.exceptions import ArraySizeError
 from dissect.cstruct.expression import Expression
 
 if TYPE_CHECKING:
+    from typing_extensions import Self
+
     from dissect.cstruct.cstruct import cstruct
 
 
@@ -27,10 +29,10 @@ class MetaType(type):
     """The alignment of the type in bytes. A value of ``None`` will be treated as 1-byte aligned."""
 
     # This must be the actual type, but since Array is a subclass of BaseType, we correct this at the bottom of the file
-    ArrayType: type[Array] = "Array"
+    ArrayType: type[BaseArray] = "Array"
     """The array type for this type class."""
 
-    def __call__(cls, *args, **kwargs) -> MetaType | BaseType:
+    def __call__(cls, *args, **kwargs) -> Self:  # type: ignore
         """Adds support for ``TypeClass(bytes | file-like object)`` parsing syntax."""
         # TODO: add support for Type(cs) API to create new bounded type classes, similar to the old API?
         if len(args) == 1 and not isinstance(args[0], cls):
@@ -48,22 +50,26 @@ class MetaType(type):
 
         return type.__call__(cls, *args, **kwargs)
 
-    def __getitem__(cls, num_entries: int | Expression | None) -> ArrayMetaType:
+    def __getitem__(cls, num_entries: int | Expression | None) -> type[BaseArray]:
         """Create a new array with the given number of entries."""
         return cls.cs._make_array(cls, num_entries)
 
     def __len__(cls) -> int:
         """Return the byte size of the type."""
+        # Python 3.9 compat thing for bound type vars
+        if cls is BaseType:
+            return 0
+
         if cls.size is None:
             raise TypeError("Dynamic size")
 
         return cls.size
 
-    def __default__(cls) -> BaseType:
+    def __default__(cls) -> Self:  # type: ignore
         """Return the default value of this type."""
         return cls()
 
-    def reads(cls, data: bytes) -> BaseType:
+    def reads(cls, data: bytes | memoryview | bytearray) -> Self:  # type: ignore
         """Parse the given data from a bytes-like object.
 
         Args:
@@ -74,7 +80,7 @@ class MetaType(type):
         """
         return cls._read(BytesIO(data))
 
-    def read(cls, obj: BinaryIO | bytes) -> BaseType:
+    def read(cls, obj: BinaryIO | bytes | memoryview | bytearray) -> Self:  # type: ignore
         """Parse the given data.
 
         Args:
@@ -86,6 +92,9 @@ class MetaType(type):
         if _is_buffer_type(obj):
             return cls.reads(obj)
 
+        if not _is_readable_type(obj):
+            raise TypeError("Invalid object type")
+
         return cls._read(obj)
 
     def write(cls, stream: BinaryIO, value: Any) -> int:
@@ -113,7 +122,7 @@ class MetaType(type):
         cls._write(out, value)
         return out.getvalue()
 
-    def _read(cls, stream: BinaryIO, context: dict[str, Any] | None = None) -> BaseType:
+    def _read(cls, stream: BinaryIO, context: dict[str, Any] | None = None) -> Self:  # type: ignore
         """Internal function for reading value.
 
         Must be implemented per type.
@@ -124,7 +133,7 @@ class MetaType(type):
         """
         raise NotImplementedError
 
-    def _read_array(cls, stream: BinaryIO, count: int, context: dict[str, Any] | None = None) -> list[BaseType]:
+    def _read_array(cls, stream: BinaryIO, count: int, context: dict[str, Any] | None = None) -> list[Self]:  # type: ignore
         """Internal function for reading array values.
 
         Allows type implementations to do optimized reading for their type.
@@ -142,7 +151,7 @@ class MetaType(type):
 
         return [cls._read(stream, context) for _ in range(count)]
 
-    def _read_0(cls, stream: BinaryIO, context: dict[str, Any] | None = None) -> list[BaseType]:
+    def _read_0(cls, stream: BinaryIO, context: dict[str, Any] | None = None) -> list[Self]:
         """Internal function for reading null-terminated data.
 
         "Null" is type specific, so must be implemented per type.
@@ -156,7 +165,7 @@ class MetaType(type):
     def _write(cls, stream: BinaryIO, data: Any) -> int:
         raise NotImplementedError
 
-    def _write_array(cls, stream: BinaryIO, array: list[BaseType]) -> int:
+    def _write_array(cls, stream: BinaryIO, array: list[Self]) -> int:  # type: ignore
         """Internal function for writing arrays.
 
         Allows type implementations to do optimized writing for their type.
@@ -167,7 +176,7 @@ class MetaType(type):
         """
         return sum(cls._write(stream, entry) for entry in array)
 
-    def _write_0(cls, stream: BinaryIO, array: list[BaseType]) -> int:
+    def _write_0(cls, stream: BinaryIO, array: list[Self]) -> int:  # type: ignore
         """Internal function for writing null-terminated arrays.
 
         Allows type implementations to do optimized writing for their type.
@@ -191,10 +200,10 @@ class _overload:
         b'\\x7b\\x00\\x00\\x00'
     """
 
-    def __init__(self, func: Callable[[Any], Any]) -> None:
+    def __init__(self, func: Callable[..., Any]) -> None:
         self.func = func
 
-    def __get__(self, instance: BaseType | None, owner: MetaType) -> Callable[[Any], bytes]:
+    def __get__(self, instance: BaseType | None, owner: type[BaseType]) -> Callable[[], bytes]:
         if instance is None:
             return functools.partial(self.func, owner)
         return functools.partial(self.func, instance.__class__, value=instance)
@@ -214,19 +223,32 @@ class BaseType(metaclass=MetaType):
         return self.__class__.size
 
 
-class ArrayMetaType(MetaType):
-    """Base metaclass for array-like types."""
+T = TypeVar("T", bound=BaseType)
+
 
-    type: MetaType
-    num_entries: int | Expression | None
-    null_terminated: bool
+class BaseArray(BaseType):
+    """Implements a fixed or dynamically sized array type.
+
+    Example:
+        When using the default C-style parser, the following syntax is supported:
 
+            x[3] -> 3 -> static length.
+            x[] -> None -> null-terminated.
+            x[expr] -> expr -> dynamic length.
+    """
+
+    type: ClassVar[type[BaseType]]
+    num_entries: ClassVar[int | Expression | None]
+    null_terminated: ClassVar[bool]
+
+    @classmethod
     def __default__(cls) -> BaseType:
         return type.__call__(
             cls, [cls.type.__default__()] * (cls.num_entries if isinstance(cls.num_entries, int) else 0)
         )
 
-    def _read(cls, stream: BinaryIO, context: dict[str, Any] | None = None) -> Array:
+    @classmethod
+    def _read(cls, stream: BinaryIO, context: dict[str, Any] | None = None) -> list[BaseType]:
         if cls.null_terminated:
             return cls.type._read_0(stream, context)
 
@@ -244,22 +266,6 @@ class ArrayMetaType(MetaType):
 
         return cls.type._read_array(stream, num, context)
 
-
-class Array(list, BaseType, metaclass=ArrayMetaType):
-    """Implements a fixed or dynamically sized array type.
-
-    Example:
-        When using the default C-style parser, the following syntax is supported:
-
-            x[3] -> 3 -> static length.
-            x[] -> None -> null-terminated.
-            x[expr] -> expr -> dynamic length.
-    """
-
-    @classmethod
-    def _read(cls, stream: BinaryIO, context: dict[str, Any] | None = None) -> Array:
-        return cls(ArrayMetaType._read(cls, stream, context))
-
     @classmethod
     def _write(cls, stream: BinaryIO, data: list[Any]) -> int:
         if cls.null_terminated:
@@ -271,11 +277,17 @@ class Array(list, BaseType, metaclass=ArrayMetaType):
         return cls.type._write_array(stream, data)
 
 
-def _is_readable_type(value: Any) -> bool:
+class Array(list[T], BaseArray):
+    @classmethod
+    def _read(cls, stream: BinaryIO, context: dict[str, Any] | None = None) -> list[T]:
+        return cls(super()._read(stream, context))
+
+
+def _is_readable_type(value: object) -> bool:
     return hasattr(value, "read")
 
 
-def _is_buffer_type(value: Any) -> bool:
+def _is_buffer_type(value: object) -> bool:
     return isinstance(value, (bytes, memoryview, bytearray))
 
 

dissect/cstruct/types/char.py

@@ -1,16 +1,23 @@
 from __future__ import annotations
 
-from typing import Any, BinaryIO
+from typing import TYPE_CHECKING, Any, BinaryIO
 
-from dissect.cstruct.types.base import EOF, ArrayMetaType, BaseType
+from dissect.cstruct.types.base import EOF, BaseArray, BaseType
 
+if TYPE_CHECKING:
+    from typing_extensions import Self
 
-class CharArray(bytes, BaseType, metaclass=ArrayMetaType):
+
+class CharArray(bytes, BaseArray):
     """Character array type for reading and writing byte strings."""
 
     @classmethod
-    def _read(cls, stream: BinaryIO, context: dict[str, Any] | None = None) -> CharArray:
-        return type.__call__(cls, ArrayMetaType._read(cls, stream, context))
+    def __default__(cls) -> Self:
+        return type.__call__(cls, b"\x00" * (0 if cls.dynamic or cls.null_terminated else cls.num_entries))
+
+    @classmethod
+    def _read(cls, stream: BinaryIO, context: dict[str, Any] | None = None) -> Self:
+        return type.__call__(cls, super()._read(stream, context))
 
     @classmethod
     def _write(cls, stream: BinaryIO, data: bytes) -> int:
@@ -24,10 +31,6 @@ class CharArray(bytes, BaseType, metaclass=ArrayMetaType):
             return stream.write(data + b"\x00")
         return stream.write(data)
 
-    @classmethod
-    def __default__(cls) -> CharArray:
-        return type.__call__(cls, b"\x00" * (0 if cls.dynamic or cls.null_terminated else cls.num_entries))
-
 
 class Char(bytes, BaseType):
     """Character type for reading and writing bytes."""
@@ -35,11 +38,15 @@ class Char(bytes, BaseType):
     ArrayType = CharArray
 
     @classmethod
-    def _read(cls, stream: BinaryIO, context: dict[str, Any] | None = None) -> Char:
+    def __default__(cls) -> Self:
+        return type.__call__(cls, b"\x00")
+
+    @classmethod
+    def _read(cls, stream: BinaryIO, context: dict[str, Any] | None = None) -> Self:
         return cls._read_array(stream, 1, context)
 
     @classmethod
-    def _read_array(cls, stream: BinaryIO, count: int, context: dict[str, Any] | None = None) -> Char:
+    def _read_array(cls, stream: BinaryIO, count: int, context: dict[str, Any] | None = None) -> Self:
         if count == 0:
             return type.__call__(cls, b"")
 
@@ -50,7 +57,7 @@ class Char(bytes, BaseType):
         return type.__call__(cls, data)
 
     @classmethod
-    def _read_0(cls, stream: BinaryIO, context: dict[str, Any] | None = None) -> Char:
+    def _read_0(cls, stream: BinaryIO, context: dict[str, Any] | None = None) -> Self:
         buf = []
         while True:
             byte = stream.read(1)
@@ -73,7 +80,3 @@ class Char(bytes, BaseType):
             data = data.encode("latin-1")
 
         return stream.write(data)
-
-    @classmethod
-    def __default__(cls) -> Char:
-        return type.__call__(cls, b"\x00")

dissect/cstruct/types/enum.py

@@ -1,30 +1,41 @@
 from __future__ import annotations
 
 import sys
+from enum import Enum as _Enum
 from enum import EnumMeta, IntEnum, IntFlag
-from typing import TYPE_CHECKING, Any, BinaryIO
+from typing import TYPE_CHECKING, Any, BinaryIO, TypeVar, overload
 
 from dissect.cstruct.types.base import Array, BaseType, MetaType
 
 if TYPE_CHECKING:
+    from typing_extensions import Self
+
     from dissect.cstruct.cstruct import cstruct
 
 
 PY_311 = sys.version_info >= (3, 11, 0)
 PY_312 = sys.version_info >= (3, 12, 0)
 
+_S = TypeVar("_S")
+
 
 class EnumMetaType(EnumMeta, MetaType):
-    type: MetaType
+    type: type[BaseType]
+
+    @overload
+    def __call__(cls, value: cstruct, name: str, type_: type[BaseType], *args, **kwargs) -> type[Enum]: ...
+
+    @overload
+    def __call__(cls: type[_S], value: int | BinaryIO | bytes) -> _S: ...
 
     def __call__(
         cls,
-        value: cstruct | int | BinaryIO | bytes = None,
+        value: cstruct | int | BinaryIO | bytes | None = None,
         name: str | None = None,
-        type_: MetaType | None = None,
+        type_: type[BaseType] | None = None,
         *args,
         **kwargs,
-    ) -> EnumMetaType:
+    ) -> Enum | type[Enum]:
         if name is None:
             if value is None:
                 value = cls.type.__default__()
@@ -35,6 +46,8 @@ class EnumMetaType(EnumMeta, MetaType):
 
             return super().__call__(value)
 
+        # We are constructing a new Enum class
+        # cs is the cstruct instance, but we can't isinstance check it due to circular imports
         cs = value
         if not issubclass(type_, int):
             raise TypeError("Enum can only be created from int type")
@@ -50,7 +63,13 @@ class EnumMetaType(EnumMeta, MetaType):
 
         return enum_cls
 
-    def __getitem__(cls, name: str | int) -> Enum | Array:
+    @overload
+    def __getitem__(cls: type[_S], name: str) -> _S: ...
+
+    @overload
+    def __getitem__(cls: type[_S], name: int) -> Array: ...
+
+    def __getitem__(cls: type[_S], name: str | int) -> _S | Array:
         if isinstance(name, str):
             return super().__getitem__(name)
         return MetaType.__getitem__(cls, name)
@@ -64,24 +83,24 @@ class EnumMetaType(EnumMeta, MetaType):
                 return True
             return value in cls._value2member_map_
 
-    def _read(cls, stream: BinaryIO, context: dict[str, Any] | None = None) -> Enum:
+    def _read(cls, stream: BinaryIO, context: dict[str, Any] | None = None) -> Self:
         return cls(cls.type._read(stream, context))
 
-    def _read_array(cls, stream: BinaryIO, count: int, context: dict[str, Any] | None = None) -> list[Enum]:
+    def _read_array(cls, stream: BinaryIO, count: int, context: dict[str, Any] | None = None) -> list[Self]:
         return list(map(cls, cls.type._read_array(stream, count, context)))
 
-    def _read_0(cls, stream: BinaryIO, context: dict[str, Any] | None = None) -> list[Enum]:
+    def _read_0(cls, stream: BinaryIO, context: dict[str, Any] | None = None) -> list[Self]:
         return list(map(cls, cls.type._read_0(stream, context)))
 
     def _write(cls, stream: BinaryIO, data: Enum) -> int:
         return cls.type._write(stream, data.value)
 
-    def _write_array(cls, stream: BinaryIO, array: list[Enum]) -> int:
-        data = [entry.value if isinstance(entry, Enum) else entry for entry in array]
+    def _write_array(cls, stream: BinaryIO, array: list[BaseType | int]) -> int:
+        data = [entry.value if isinstance(entry, _Enum) else entry for entry in array]
         return cls.type._write_array(stream, data)
 
-    def _write_0(cls, stream: BinaryIO, array: list[BaseType]) -> int:
-        data = [entry.value if isinstance(entry, Enum) else entry for entry in array]
+    def _write_0(cls, stream: BinaryIO, array: list[BaseType | int]) -> int:
+        data = [entry.value if isinstance(entry, _Enum) else entry for entry in array]
         return cls._write_array(stream, [*data, cls.type.__default__()])
 
 
@@ -180,7 +199,7 @@ class Enum(BaseType, IntEnum, metaclass=EnumMetaType):
         return hash((self.__class__, self.name, self.value))
 
     @classmethod
-    def _missing_(cls, value: int) -> Enum:
+    def _missing_(cls, value: int) -> Self:
         # Emulate FlagBoundary.KEEP for enum (allow values other than the defined members)
         new_member = int.__new__(cls, value)
         new_member._name_ = None

dissect/cstruct/types/int.py

@@ -1,10 +1,13 @@
 from __future__ import annotations
 
-from typing import Any, BinaryIO
+from typing import TYPE_CHECKING, Any, BinaryIO
 
 from dissect.cstruct.types.base import BaseType
 from dissect.cstruct.utils import ENDIANNESS_MAP
 
+if TYPE_CHECKING:
+    from typing_extensions import Self
+
 
 class Int(int, BaseType):
     """Integer type that can span an arbitrary amount of bytes."""
@@ -12,7 +15,7 @@ class Int(int, BaseType):
     signed: bool
 
     @classmethod
-    def _read(cls, stream: BinaryIO, context: dict[str, Any] | None = None) -> Int:
+    def _read(cls, stream: BinaryIO, context: dict[str, Any] | None = None) -> Self:
         data = stream.read(cls.size)
 
         if len(data) != cls.size:
@@ -21,7 +24,7 @@ class Int(int, BaseType):
         return cls.from_bytes(data, ENDIANNESS_MAP[cls.cs.endian], signed=cls.signed)
 
     @classmethod
-    def _read_0(cls, stream: BinaryIO, context: dict[str, Any] | None = None) -> Int:
+    def _read_0(cls, stream: BinaryIO, context: dict[str, Any] | None = None) -> Self:
         result = []
 
         while True: