API

class BaseSnapshot(abc.ABC):
    """Base class for all `*Snapshot` classes."""

    def __init__(self, type_: Snap, pid: int = 0) -> None:  # noqa
        self._type = type_
        self._pid = pid
        ptr = CreateToolhelp32Snapshot(type_, pid)
        if ptr == INVALID_HANDLE_VALUE:
            raise CreationFailedError
        self._handle = ptr

    def __enter__(self, *_) -> Self:
        return self

    def __repr__(self) -> str:
        return f"<Snapshot type={self._type}, pid={self._pid}>"

    def __exit__(self, *_) -> None:
        self.close()

    def close(self) -> bool:
        """Close the underlying snapshot handle and free up resources.

        Returns:
            bool: Whether the underlying handle was closed successfully.
        """
        if self._handle:  # pragma: no cover
            return bool(CloseHandle(self._handle))
        return False

class Snapshot(BaseSnapshot):
    """All-in-one snapshot."""

    def __init__(self, pid: int = 0, include_32bit: bool = False) -> None:
        """Create a snapshot containing the heaps, modules, processes and threads.

        Args:
            pid (int, optional): The PID of the process whose modules and heaps
                are to be enumerated. Defaults to 0. If it is 0, then the
                enumeration takes place for the current / calling process.
            include_32bit (bool, optional): Passed to `ModuleSnapshot`.

        Raises:
            CreationFailedError: If the call fails internally.
        """
        snap_type = Snap.ALL
        if include_32bit:
            snap_type |= Snap.MODULE32
        super().__init__(snap_type, pid)
        self._pid = pid
        self._he = HEAPENTRY32()
        self._hl = HEAPLIST32()
        self._me = MODULEENTRY32W()
        self._pe = PROCESSENTRY32W()
        self._te = THREADENTRY32()

    @property
    def heaps(self) -> HeapIterator:
        """Iterator for the heaps present in the snapshot."""
        return HeapIterator(self._handle, self._hl)

    @property
    def modules(self) -> ModuleIterator:
        """Iterator for the modules present in the snapshot."""
        return ModuleIterator(self._handle, self._me)

    @property
    def processes(self) -> ProcessIterator:
        """Iterator for the processes present in the snapshot."""
        return ProcessIterator(self._handle, self._pe)

    @property
    def threads(self) -> ThreadIterator:
        """Iterator for the threads present in the snapshot."""
        return ThreadIterator(self._handle, self._te)

class HeapListSnapshot(BaseSnapshot):
    """Enumerates the heaps used by a particular process.

    Enumerating (a.k.a. walking) the heap is an expensive operation, look for
    ways to avoid redoing it or use the Win32 `HeapWalk` function instead for
    a much better performance.
    """

    def __init__(self, pid: int = 0) -> None:
        """Create a module snapshot.

        Args:
            pid (int, optional): The PID of the process whose modules are to
                be enumerated. Defaults to 0. If it is 0, then the modules for
                the current/calling process are enumerated.

        Raises:
            CreationFailedError: If the call fails internally.
        """
        super().__init__(Snap.HEAPLIST, pid)
        self._hl = HEAPLIST32()

    def __iter__(self) -> HeapIterator:
        return HeapIterator(self._handle, self._hl)

class ModuleSnapshot(BaseSnapshot):
    """Enumerates the modules belonging to a particular process."""

    def __init__(self, pid: int = 0, include_32bit: bool = False) -> None:
        """Create a module snapshot.

        Args:
            pid (int, optional): The PID of the process whose modules are to
                be enumerated. Defaults to 0. If it is 0, then the modules for
                the current/calling process are enumerated.
            include_32bit (bool, optional): Enumerate 32-bit modules also.
                Valid only for 64-bit processes, else this parameter is
                ignored. Defaults to False.

        Raises:
            CreationFailedError: If the call fails internally.
        """
        snap_type = Snap.MODULE
        if include_32bit:
            snap_type |= Snap.MODULE32
        super().__init__(snap_type, pid)
        self._me = MODULEENTRY32W()

    def __iter__(self) -> ModuleIterator | None:
        return ModuleIterator(self._handle, self._me)

class ProcessSnapshot(BaseSnapshot):
    """Enumerates all the processes residing in the system address space."""

    def __init__(self) -> None:
        """Create a process snapshot.

        Raises:
            CreationFailedError: If the call fails internally.
        """
        super().__init__(Snap.PROCESS)
        self._pe = PROCESSENTRY32W()

    def __iter__(self) -> ProcessIterator | None:
        return ProcessIterator(self._handle, self._pe)

class ThreadSnapshot(BaseSnapshot):
    """Enumerates all the threads executing on the system."""

    def __init__(self) -> None:
        """Create a thread snapshot.

        Raises:
            CreationFailedError: If the call fails internally.
        """
        super().__init__(Snap.THREAD)
        self._te = THREADENTRY32()

    def __iter__(self) -> ThreadIterator | None:
        return ThreadIterator(self._handle, self._te)