iheap 0.2.0__py3-none-any.whl

iheap/__init__.py

@@ -0,0 +1,36 @@
+"""Indexed priority queues with a Rust core.
+
+* :class:`IndexedHeap` -- low-level heap of float priorities addressed by
+  integer handles. Holds no Python objects.
+* :class:`ObjectHeap` -- the same, for any comparable priorities (strings,
+  tuples, ...), in pure Python.
+* :class:`PQDict` -- ``pqdict``-style mapping of keys to values ordered by
+  priority. Keys and values live in Python; only priorities reach the core.
+
+:data:`NATIVE` tells whether the Rust core is in use. Where no prebuilt wheel
+exists for the interpreter (for example free-threaded CPython 3.13), a
+pure-Python core with the same API is used instead.
+"""
+
+from importlib.metadata import PackageNotFoundError
+from importlib.metadata import version as _version
+
+from ._core import NATIVE, IndexedHeap
+from ._objheap import ObjectHeap
+from ._pqdict import PQDict, maxpq, minpq, nlargest, nsmallest
+
+__all__ = [
+    "NATIVE",
+    "IndexedHeap",
+    "ObjectHeap",
+    "PQDict",
+    "maxpq",
+    "minpq",
+    "nlargest",
+    "nsmallest",
+]
+
+try:
+    __version__ = _version("iheap")
+except PackageNotFoundError:  # running from a source checkout
+    __version__ = "0+unknown"

iheap/_core.py

@@ -0,0 +1,21 @@
+"""Selects the ``IndexedHeap`` implementation.
+
+The native extension (Rust) is used whenever it can be imported. Otherwise,
+or when the environment variable ``IHEAP_PURE_PYTHON`` is set to a non-empty
+value, the pure-Python implementation with the same API is used.
+"""
+
+import os
+
+from ._pyheap import IndexedHeap
+
+NATIVE = False
+if not os.environ.get("IHEAP_PURE_PYTHON"):
+    try:
+        from ._iheap import IndexedHeap  # type: ignore[assignment]  # noqa: F811
+    except ImportError:  # no native wheel for this interpreter, e.g. CPython 3.13t
+        pass
+    else:
+        NATIVE = True
+
+__all__ = ["NATIVE", "IndexedHeap"]

iheap/_iheap.pyi

@@ -0,0 +1,155 @@
+from typing import Any, final
+
+@final
+class IndexedHeap:
+    """Native indexed priority queue of float priorities addressed by handles.
+
+    [`push`][iheap.IndexedHeap.push] returns an integer *handle* that stays
+    valid until its entry is popped or removed, and lets you update or remove
+    that entry in ``O(log n)``. Handles are small non-negative integers below
+    [`handle_bound`][iheap.IndexedHeap.handle_bound] and are recycled, so you
+    can keep per-entry data in a plain list indexed by handle. The heap
+    itself holds no Python objects.
+
+    Priorities must be convertible to ``float`` and not NaN. Integers are
+    exact up to 2**53.
+
+    Thread safety: a heap may be shared between threads, including on
+    free-threaded Python. Each call runs under an internal lock, so
+    concurrent calls are serialized; sequences of calls are not atomic.
+
+    Args:
+        reverse: If true, the largest priority is on top (max-heap).
+        capacity: Number of entries to pre-allocate room for.
+
+    Examples:
+        >>> heap = IndexedHeap()
+        >>> a, b = heap.push(3.0), heap.push(1.0)
+        >>> heap.peek() == (b, 1.0)
+        True
+        >>> heap.update(a, 0.5)
+        3.0
+        >>> heap.pop() == (a, 0.5)
+        True
+    """
+
+    def __init__(self, *, reverse: bool = False, capacity: int = 0) -> None: ...
+    @property
+    def reverse(self) -> bool:
+        """``True`` for a max-heap."""
+    @property
+    def handle_bound(self) -> int:
+        """Exclusive upper bound of every live handle (the handle high-water mark)."""
+    def __len__(self) -> int:
+        """Number of live entries."""
+    def __bool__(self) -> bool: ...
+    def __contains__(self, handle: int) -> bool:
+        """Whether ``handle`` refers to a live entry."""
+    def push(self, priority: float) -> int:
+        """Insert an entry in ``O(log n)``.
+
+        Args:
+            priority: Priority of the new entry.
+
+        Returns:
+            The handle of the new entry.
+
+        Raises:
+            ValueError: ``priority`` is NaN.
+            TypeError: ``priority`` is not convertible to ``float``.
+        """
+    def push_many(self, priorities: list[float]) -> list[int]:
+        """Insert many entries at once.
+
+        Large batches are appended and re-heapified in ``O(n)``. Atomic: if any
+        priority is invalid, nothing is inserted. On an empty heap after
+        construction or [`clear`][iheap.IndexedHeap.clear], the returned
+        handles are ``0, 1, ..., len(priorities) - 1``.
+
+        Args:
+            priorities: Priorities of the new entries.
+
+        Returns:
+            Handles of the new entries, in input order.
+
+        Raises:
+            ValueError: A priority is NaN.
+        """
+    def peek(self) -> tuple[int, float] | None:
+        """Return ``(handle, priority)`` of the top entry, or ``None`` if empty."""
+    def pop(self) -> tuple[int, float]:
+        """Remove and return the top entry in ``O(log n)``.
+
+        Returns:
+            ``(handle, priority)`` of the removed entry. The handle becomes
+            invalid and may be reused by later pushes.
+
+        Raises:
+            IndexError: The heap is empty.
+        """
+    def push_pop(self, priority: float) -> tuple[int, float] | None:
+        """Fused push-then-pop in a single pass (like ``heapq.heappushpop``).
+
+        Args:
+            priority: Priority of the new entry.
+
+        Returns:
+            ``None`` if the new entry would be popped right away (the heap is
+            empty or its top does not strictly precede ``priority``); the heap
+            is then unchanged. Otherwise the top entry is popped, **its handle
+            is reassigned to the new entry**, and ``(handle, popped_priority)``
+            is returned.
+
+        Raises:
+            ValueError: ``priority`` is NaN.
+        """
+    def pop_many(self, n: int) -> list[tuple[int, float]]:
+        """Pop up to ``n`` entries.
+
+        Returns:
+            ``(handle, priority)`` pairs in priority order.
+        """
+    def priority(self, handle: int) -> float:
+        """Return the priority of ``handle``.
+
+        Raises:
+            KeyError: ``handle`` is not live.
+        """
+    def update(self, handle: int, priority: float) -> float:
+        """Change the priority of ``handle`` in ``O(log n)``.
+
+        Returns:
+            The previous priority.
+
+        Raises:
+            KeyError: ``handle`` is not live.
+            ValueError: ``priority`` is NaN.
+        """
+    def update_many(self, handles: list[int], priorities: list[float]) -> None:
+        """Change many priorities at once.
+
+        Re-heapifies in ``O(n)`` when that is cheaper than individual updates.
+        Atomic: on error nothing changes. If a handle appears more than once,
+        the last priority wins.
+
+        Raises:
+            KeyError: A handle is not live.
+            ValueError: The lists differ in length, or a priority is NaN.
+        """
+    def remove(self, handle: int) -> float:
+        """Remove ``handle`` in ``O(log n)``.
+
+        Returns:
+            The priority of the removed entry.
+
+        Raises:
+            KeyError: ``handle`` is not live.
+        """
+    def clear(self) -> None:
+        """Remove every entry. Handle numbering restarts at 0."""
+    def items(self) -> list[tuple[int, float]]:
+        """Return live ``(handle, priority)`` pairs in handle order (not priority order)."""
+    def copy(self) -> IndexedHeap:
+        """Return an independent copy, handles included."""
+    def __copy__(self) -> IndexedHeap: ...
+    def __deepcopy__(self, memo: Any) -> IndexedHeap: ...

iheap/_objheap.py

@@ -0,0 +1,80 @@
+"""``ObjectHeap``: an indexed heap over arbitrary comparable priorities.
+
+The native and float cores only order numbers. ``ObjectHeap`` orders any
+priorities that support ``<`` among themselves: strings, tuples (e.g.
+``(priority, sequence)`` for FIFO ties), dates, ``Decimal`` exactly, and so
+on. It reuses the pure-Python heap, which only ever compares keys with ``<``
+and ``>=``; a max-heap wraps keys in :class:`_Reversed`, which inverts both.
+It runs at pure-Python speed.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from ._pyheap import IndexedHeap as _PyHeap
+
+
+class _Reversed:
+    """Key wrapper that inverts the order of the wrapped priority."""
+
+    __slots__ = ("value",)
+
+    def __init__(self, value: Any) -> None:
+        self.value = value
+
+    def __lt__(self, other: _Reversed) -> bool:
+        return other.value < self.value
+
+    def __ge__(self, other: _Reversed) -> bool:
+        return not other.value < self.value
+
+
+class ObjectHeap(_PyHeap):
+    """Indexed heap whose priorities are any mutually comparable objects.
+
+    Same API, handle semantics and thread safety as
+    [`IndexedHeap`][iheap.IndexedHeap], but priorities are stored as given
+    (not converted to ``float``) and ordered with ``<``. They must be totally
+    ordered and comparable with each other.
+
+    Raises:
+        TypeError: A new priority cannot be compared with those in the heap
+            (checked against the current top before anything changes).
+        ValueError: A priority is NaN-like (``p != p``).
+
+    Examples:
+        >>> heap = ObjectHeap()
+        >>> h = heap.push("banana")
+        >>> _ = heap.push_many(["cherry", "apple"])
+        >>> heap.pop()[1]
+        'apple'
+        >>> heap.update(h, "zucchini")
+        'banana'
+        >>> heap.peek()[1]
+        'cherry'
+    """
+
+    __slots__ = ()
+
+    def _key(self, priority: Any) -> Any:
+        if priority != priority:
+            raise ValueError("priority must not be NaN")
+        return _Reversed(priority) if self._sign < 0 else priority
+
+    def _decode(self, key: Any) -> Any:
+        return key.value if self._sign < 0 else key
+
+    def _check(self, keys: list[Any]) -> None:
+        # Mixing incomparable types (say `str` and `int`) raises inside a
+        # sift, half-way through moving entries. Comparing every new key with
+        # a reference first (the current top, else the first new key) makes
+        # such operations fail with no effect.
+        if self._heap:
+            reference = self._keys[self._heap[0]]
+        elif keys:
+            reference = keys[0]
+        else:
+            return
+        for key in keys:
+            bool(key < reference)