fastquadtree 0.6.1__cp38-abi3-win_amd64.whl → 0.8.0__cp38-abi3-win_amd64.whl

fastquadtree/_bimap.py

@@ -1,14 +1,16 @@
 # _bimap.py
 from __future__ import annotations
 
-from typing import Any, Iterable, Iterator
+from typing import Any, Generic, Iterable, Iterator, TypeVar
 
-from ._item import Item
+from ._item import Item  # base class for PointItem and RectItem
 
+TItem = TypeVar("TItem", bound=Item)
 
-class BiMap:
+
+class BiMap(Generic[TItem]):
     """
-    Bidirectional map to the same Item:
+    Bidirectional map to the same Item subtype:
       id -> Item
       obj -> Item  (uses object identity)
 
@@ -20,19 +22,22 @@ class BiMap:
 
     __slots__ = ("_id_to_item", "_objid_to_item")
 
-    def __init__(self, items: Iterable[Item] | None = None) -> None:
-        self._id_to_item: dict[int, Item] = {}
-        self._objid_to_item: dict[int, Item] = {}
+    def __init__(
+        self,
+        items: Iterable[TItem] | None = None,
+    ) -> None:
+        self._id_to_item: dict[int, TItem] = {}
+        self._objid_to_item: dict[int, TItem] = {}
         if items:
             for it in items:
                 self.add(it)
 
     # - core -
 
-    def add(self, item: Item) -> None:
+    def add(self, item: TItem) -> None:
         """
         Insert or replace mapping for this Item.
-        Handles conflicts so that both id and obj point to this exact Item.
+        Handles conflicts so both id and obj point to this exact Item.
         """
         id_ = item.id_
         obj = item.obj
@@ -55,13 +60,13 @@ class BiMap:
         if obj is not None:
             self._objid_to_item[id(obj)] = item
 
-    def by_id(self, id_: int) -> Item | None:
+    def by_id(self, id_: int) -> TItem | None:
         return self._id_to_item.get(id_)
 
-    def by_obj(self, obj: Any) -> Item | None:
+    def by_obj(self, obj: Any) -> TItem | None:
         return self._objid_to_item.get(id(obj))
 
-    def pop_id(self, id_: int) -> Item | None:
+    def pop_id(self, id_: int) -> TItem | None:
         it = self._id_to_item.pop(id_, None)
         if it is not None:
             obj = it.obj
@@ -69,25 +74,20 @@ class BiMap:
                 self._objid_to_item.pop(id(obj), None)
         return it
 
-    def pop_obj(self, obj: Any) -> Item | None:
+    def pop_obj(self, obj: Any) -> TItem | None:
         it = self._objid_to_item.pop(id(obj), None)
         if it is not None:
             self._id_to_item.pop(it.id_, None)
         return it
 
-    def pop_item(self, item: Item) -> Item | None:
+    def pop_item(self, item: TItem) -> TItem | None:
         """
         Remove this exact Item if present on either side.
         """
-        removed = None
-        # Remove by id first
-        removed = self._id_to_item.pop(item.id_)
-
-        # Remove by obj side
+        removed = self._id_to_item.pop(item.id_, None)
         obj = item.obj
         if obj is not None:
             self._objid_to_item.pop(id(obj), None)
-            removed = removed or item
         return removed
 
     # - convenience -
@@ -105,8 +105,8 @@ class BiMap:
     def contains_obj(self, obj: Any) -> bool:
         return id(obj) in self._objid_to_item
 
-    def items_by_id(self) -> Iterator[tuple[int, Item]]:
+    def items_by_id(self) -> Iterator[tuple[int, TItem]]:
         return iter(self._id_to_item.items())
 
-    def items(self) -> Iterator[Item]:
+    def items(self) -> Iterator[TItem]:
         return iter(self._id_to_item.values())

fastquadtree/_item.py

@@ -1,7 +1,13 @@
 # item.py
 from __future__ import annotations
 
-from typing import Any
+from typing import Any, Tuple
+
+Bounds = Tuple[float, float, float, float]
+"""Axis-aligned rectangle as (min_x, min_y, max_x, max_y)."""
+
+Point = Tuple[float, float]
+"""2D point as (x, y)."""
 
 
 class Item:
@@ -10,15 +16,37 @@ class Item:
 
     Attributes:
         id_: Integer identifier.
-        x: X coordinate.
-        y: Y coordinate.
+        geom: The geometry, either a Point or Rectangle Bounds.
         obj: The attached Python object if available, else None.
     """
 
-    __slots__ = ("id_", "obj", "x", "y")
+    __slots__ = ("geom", "id_", "obj")
 
-    def __init__(self, id_: int, x: float, y: float, obj: Any | None = None):
+    def __init__(self, id_: int, geom: Point | Bounds, obj: Any | None = None):
         self.id_ = id_
-        self.x = x
-        self.y = y
+        self.geom = geom
         self.obj = obj
+
+
+class PointItem(Item):
+    """
+    Lightweight point item wrapper for tracking and as_items results.
+    """
+
+    __slots__ = ("geom", "id_", "obj", "x", "y")
+
+    def __init__(self, id_: int, geom: Point, obj: Any | None = None):
+        super().__init__(id_, geom, obj)
+        self.x, self.y = geom
+
+
+class RectItem(Item):
+    """
+    Lightweight rectangle item wrapper for tracking and as_items results.
+    """
+
+    __slots__ = ("geom", "id_", "max_x", "max_y", "min_x", "min_y", "obj")
+
+    def __init__(self, id_: int, geom: Bounds, obj: Any | None = None):
+        super().__init__(id_, geom, obj)
+        self.min_x, self.min_y, self.max_x, self.max_y = geom

fastquadtree/point_quadtree.py

@@ -0,0 +1,161 @@
+# point_quadtree.py
+from __future__ import annotations
+
+from typing import Any, Literal, Tuple, overload
+
+from ._base_quadtree import Bounds, _BaseQuadTree
+from ._item import Point, PointItem
+from ._native import QuadTree as _RustQuadTree  # native point tree
+
+_IdCoord = Tuple[int, float, float]
+
+
+class QuadTree(_BaseQuadTree[Point, _IdCoord, PointItem]):
+    """
+    Point version of the quadtree. All geometries are 2D points (x, y).
+    High-level Python wrapper over the Rust quadtree engine.
+
+    Performance characteristics:
+        Inserts: average O(log n) <br>
+        Rect queries: average O(log n + k) where k is matches returned <br>
+        Nearest neighbor: average O(log n) <br>
+
+    Thread-safety:
+        Instances are not thread-safe. Use external synchronization if you
+        mutate the same tree from multiple threads.
+
+    Args:
+        bounds: World bounds as (min_x, min_y, max_x, max_y).
+        capacity: Max number of points per node before splitting.
+        max_depth: Optional max tree depth. If omitted, engine decides.
+        track_objects: Enable id <-> object mapping inside Python.
+        start_id: Starting auto-assigned id when you omit id on insert.
+
+    Raises:
+        ValueError: If parameters are invalid or inserts are out of bounds.
+    """
+
+    def __init__(
+        self,
+        bounds: Bounds,
+        capacity: int,
+        *,
+        max_depth: int | None = None,
+        track_objects: bool = False,
+        start_id: int = 1,
+    ):
+        super().__init__(
+            bounds,
+            capacity,
+            max_depth=max_depth,
+            track_objects=track_objects,
+            start_id=start_id,
+        )
+
+    @overload
+    def query(
+        self, rect: Bounds, *, as_items: Literal[False] = ...
+    ) -> list[_IdCoord]: ...
+    @overload
+    def query(self, rect: Bounds, *, as_items: Literal[True]) -> list[PointItem]: ...
+    def query(
+        self, rect: Bounds, *, as_items: bool = False
+    ) -> list[PointItem] | list[_IdCoord]:
+        """
+        Return all points inside an axis-aligned rectangle.
+
+        Args:
+            rect: Query rectangle as (min_x, min_y, max_x, max_y).
+            as_items: If True, return Item wrappers. If False, return raw tuples.
+
+        Returns:
+            If as_items is False: list of (id, x, y) tuples.
+            If as_items is True: list of Item objects.
+        """
+        raw = self._native.query(rect)
+        if not as_items:
+            return raw
+        if self._items is None:
+            raise ValueError("Cannot return results as items with track_objects=False")
+        out: list[PointItem] = []
+        for id_, _x, _y in raw:
+            it = self._items.by_id(id_)
+            if it is None:
+                raise RuntimeError(
+                    f"Internal error: id {id_} present in native tree but missing from tracker."
+                )
+            out.append(it)
+        return out
+
+    @overload
+    def nearest_neighbor(
+        self, xy: Point, *, as_item: Literal[False] = ...
+    ) -> _IdCoord | None: ...
+    @overload
+    def nearest_neighbor(
+        self, xy: Point, *, as_item: Literal[True]
+    ) -> PointItem | None: ...
+    def nearest_neighbor(
+        self, xy: Point, *, as_item: bool = False
+    ) -> PointItem | _IdCoord | None:
+        """
+        Return the single nearest neighbor to the query point.
+
+        Args:
+            xy: Query point (x, y).
+            as_item: If True, return Item. If False, return (id, x, y).
+
+        Returns:
+            The nearest neighbor or None if the tree is empty.
+        """
+        t = self._native.nearest_neighbor(xy)
+        if t is None or not as_item:
+            return t
+        if self._items is None:
+            raise ValueError("Cannot return result as item with track_objects=False")
+        id_, _x, _y = t
+        it = self._items.by_id(id_)
+        if it is None:
+            raise RuntimeError("Internal error: missing tracked item")
+        return it
+
+    @overload
+    def nearest_neighbors(
+        self, xy: Point, k: int, *, as_items: Literal[False] = ...
+    ) -> list[_IdCoord]: ...
+    @overload
+    def nearest_neighbors(
+        self, xy: Point, k: int, *, as_items: Literal[True]
+    ) -> list[PointItem]: ...
+    def nearest_neighbors(self, xy: Point, k: int, *, as_items: bool = False):
+        """
+        Return the k nearest neighbors to the query point in order of increasing distance.
+
+        Args:
+            xy: Query point (x, y).
+            k: Number of neighbors to return.
+            as_items: If True, return Item wrappers. If False, return raw tuples.
+        Returns:
+            If as_items is False: list of (id, x, y) tuples.
+            If as_items is True: list of Item objects.
+        """
+        raw = self._native.nearest_neighbors(xy, k)
+        if not as_items:
+            return raw
+        if self._items is None:
+            raise ValueError("Cannot return results as items with track_objects=False")
+        out: list[PointItem] = []
+        for id_, _x, _y in raw:
+            it = self._items.by_id(id_)
+            if it is None:
+                raise RuntimeError("Internal error: missing tracked item")
+            out.append(it)
+        return out
+
+    def _new_native(self, bounds: Bounds, capacity: int, max_depth: int | None) -> Any:
+        if max_depth is None:
+            return _RustQuadTree(bounds, capacity)
+        return _RustQuadTree(bounds, capacity, max_depth=max_depth)
+
+    def _make_item(self, id_: int, geom: Point, obj: Any | None) -> PointItem:
+        return PointItem(id_, geom, obj)

fastquadtree/rect_quadtree.py

@@ -0,0 +1,98 @@
+# rect_quadtree.py
+from __future__ import annotations
+
+from typing import Any, Literal, Tuple, overload
+
+from ._base_quadtree import Bounds, _BaseQuadTree
+from ._item import RectItem
+from ._native import RectQuadTree as _RustRectQuadTree  # native rect tree
+
+_IdRect = Tuple[int, float, float, float, float]
+Point = Tuple[float, float]  # only for type hints in docstrings
+
+
+class RectQuadTree(_BaseQuadTree[Bounds, _IdRect, RectItem]):
+    """
+    Rectangle version of the quadtree. All geometries are axis-aligned rectangles. (min_x, min_y, max_x, max_y)
+    High-level Python wrapper over the Rust quadtree engine.
+
+    Performance characteristics:
+        Inserts: average O(log n) <br>
+        Rect queries: average O(log n + k) where k is matches returned <br>
+        Nearest neighbor: average O(log n) <br>
+
+    Thread-safety:
+        Instances are not thread-safe. Use external synchronization if you
+        mutate the same tree from multiple threads.
+
+    Args:
+        bounds: World bounds as (min_x, min_y, max_x, max_y).
+        capacity: Max number of points per node before splitting.
+        max_depth: Optional max tree depth. If omitted, engine decides.
+        track_objects: Enable id <-> object mapping inside Python.
+        start_id: Starting auto-assigned id when you omit id on insert.
+
+    Raises:
+        ValueError: If parameters are invalid or inserts are out of bounds.
+    """
+
+    def __init__(
+        self,
+        bounds: Bounds,
+        capacity: int,
+        *,
+        max_depth: int | None = None,
+        track_objects: bool = False,
+        start_id: int = 1,
+    ):
+        super().__init__(
+            bounds,
+            capacity,
+            max_depth=max_depth,
+            track_objects=track_objects,
+            start_id=start_id,
+        )
+
+    @overload
+    def query(
+        self, rect: Bounds, *, as_items: Literal[False] = ...
+    ) -> list[_IdRect]: ...
+    @overload
+    def query(self, rect: Bounds, *, as_items: Literal[True]) -> list[RectItem]: ...
+    def query(self, rect: Bounds, *, as_items: bool = False):
+        """
+        Query the tree for all items that intersect the given rectangle.
+
+        Args:
+            rect: Query rectangle as (min_x, min_y, max_x, max_y).
+            as_items: If True, return Item wrappers. If False, return raw tuples.
+
+        Returns:
+            If as_items is False: list of (id, x0, y0, x1, y1) tuples.
+            If as_items is True: list of Item objects.
+        """
+        raw = self._native.query(rect)
+        if not as_items:
+            return raw
+        if self._items is None:
+            # Build RectItem without objects
+            return [
+                RectItem(id_, (x0, y0, x1, y1), None) for (id_, x0, y0, x1, y1) in raw
+            ]
+        out: list[RectItem] = []
+        for id_, _x0, _y0, _x1, _y1 in raw:
+            it = self._items.by_id(id_)
+            if it is None:
+                raise RuntimeError(
+                    f"Internal error: id {id_} present in native tree but missing from tracker."
+                )
+            out.append(it)
+        return out
+
+    def _new_native(self, bounds: Bounds, capacity: int, max_depth: int | None) -> Any:
+        if max_depth is None:
+            return _RustRectQuadTree(bounds, capacity)
+        return _RustRectQuadTree(bounds, capacity, max_depth=max_depth)
+
+    def _make_item(self, id_: int, geom: Bounds, obj: Any | None) -> RectItem:
+        return RectItem(id_, geom, obj)