fastquadtree 0.9.1__tar.gz → 1.0.1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (73) hide show
  1. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/Cargo.lock +1 -1
  2. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/Cargo.toml +1 -1
  3. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/PKG-INFO +3 -1
  4. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/README.md +2 -0
  5. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/benchmarks/benchmark_native_vs_shim.py +23 -8
  6. fastquadtree-1.0.1/docs/api/pyqtree.md +4 -0
  7. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/docs/benchmark.md +25 -6
  8. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/docs/index.md +3 -1
  9. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/docs/quickstart.md +5 -5
  10. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/mkdocs.yml +5 -4
  11. fastquadtree-1.0.1/pysrc/fastquadtree/_base_quadtree.py +258 -0
  12. fastquadtree-1.0.1/pysrc/fastquadtree/_obj_store.py +167 -0
  13. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/pysrc/fastquadtree/point_quadtree.py +7 -19
  14. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/pysrc/fastquadtree/pyqtree.py +48 -17
  15. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/pysrc/fastquadtree/rect_quadtree.py +4 -19
  16. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/src/lib.rs +16 -0
  17. fastquadtree-1.0.1/tests/test_base_quadtree.py +295 -0
  18. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/tests/test_clear.py +7 -8
  19. fastquadtree-1.0.1/tests/test_obj_store.py +167 -0
  20. fastquadtree-1.0.1/tests/test_point_quadtree_nn_runtime.py +46 -0
  21. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/tests/test_pyqtree_shim_compat.py +94 -0
  22. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/tests/test_python.py +24 -24
  23. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/tests/test_rect_quadtree.py +29 -25
  24. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/tests/test_wrapper_edges.py +3 -63
  25. fastquadtree-0.9.1/pysrc/fastquadtree/_base_quadtree.py +0 -263
  26. fastquadtree-0.9.1/pysrc/fastquadtree/_bimap.py +0 -112
  27. fastquadtree-0.9.1/tests/test_bimap.py +0 -155
  28. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/.github/workflows/docs.yml +0 -0
  29. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/.github/workflows/release.yml +0 -0
  30. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/.gitignore +0 -0
  31. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/.pre-commit-config.yaml +0 -0
  32. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/LICENSE +0 -0
  33. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/assets/ballpit.png +0 -0
  34. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/assets/interactive_v2_rect_screenshot.png +0 -0
  35. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/assets/interactive_v2_screenshot.png +0 -0
  36. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/assets/quadtree_bench_throughput.png +0 -0
  37. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/assets/quadtree_bench_time.png +0 -0
  38. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/benchmarks/cross_library_bench.py +0 -0
  39. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/benchmarks/quadtree_bench/__init__.py +0 -0
  40. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/benchmarks/quadtree_bench/engines.py +0 -0
  41. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/benchmarks/quadtree_bench/main.py +0 -0
  42. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/benchmarks/quadtree_bench/plotting.py +0 -0
  43. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/benchmarks/quadtree_bench/runner.py +0 -0
  44. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/benchmarks/requirements.txt +0 -0
  45. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/benchmarks/runner.py +0 -0
  46. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/benchmarks/system_info_collector.py +0 -0
  47. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/docs/api/point_item.md +0 -0
  48. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/docs/api/quadtree.md +0 -0
  49. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/docs/api/rect_item.md +0 -0
  50. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/docs/api/rect_quadtree.md +0 -0
  51. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/docs/runnables.md +0 -0
  52. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/interactive/ballpit.py +0 -0
  53. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/interactive/interactive.py +0 -0
  54. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/interactive/interactive_v2.py +0 -0
  55. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/interactive/interactive_v2_rect.py +0 -0
  56. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/interactive/requirements.txt +0 -0
  57. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/pyproject.toml +0 -0
  58. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/pysrc/fastquadtree/__init__.py +0 -0
  59. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/pysrc/fastquadtree/_item.py +0 -0
  60. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/pysrc/fastquadtree/py.typed +0 -0
  61. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/src/geom.rs +0 -0
  62. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/src/quadtree.rs +0 -0
  63. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/src/rect_quadtree.rs +0 -0
  64. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/tests/insertions.rs +0 -0
  65. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/tests/nearest_neighbor.rs +0 -0
  66. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/tests/query.rs +0 -0
  67. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/tests/rect_quadtree.rs +0 -0
  68. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/tests/rectangle_traversal.rs +0 -0
  69. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/tests/test_delete.rs +0 -0
  70. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/tests/test_delete_by_object.py +0 -0
  71. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/tests/test_delete_python.py +0 -0
  72. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/tests/test_unconventional_bounds.py +0 -0
  73. {fastquadtree-0.9.1 → fastquadtree-1.0.1}/tests/unconventional_bounds.rs +0 -0
@@ -10,7 +10,7 @@ checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8"
10
10
 
11
11
  [[package]]
12
12
  name = "fastquadtree"
13
- version = "0.9.1"
13
+ version = "1.0.1"
14
14
  dependencies = [
15
15
  "pyo3",
16
16
  "smallvec",
@@ -1,6 +1,6 @@
1
1
  [package]
2
2
  name = "fastquadtree"
3
- version = "0.9.1"
3
+ version = "1.0.1"
4
4
  edition = "2021"
5
5
  readme = "README.md"
6
6
 
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: fastquadtree
3
- Version: 0.9.1
3
+ Version: 1.0.1
4
4
  Classifier: Programming Language :: Python :: 3
5
5
  Classifier: Programming Language :: Python :: 3 :: Only
6
6
  Classifier: Programming Language :: Rust
@@ -69,6 +69,7 @@ Rust-optimized quadtree with a clean Python API
69
69
  - Fast KNN and range queries
70
70
  - Optional object tracking for id ↔ object mapping
71
71
  - [100% test coverage](https://codecov.io/gh/Elan456/fastquadtree) and CI on GitHub Actions
72
+ - Offers a drop-in [pyqtree shim](https://elan456.github.io/fastquadtree/benchmark/#pyqtree-drop-in-shim-performance-gains) that is 6.567x faster while keeping the same API
72
73
 
73
74
  ----
74
75
 
@@ -86,6 +87,7 @@ pip install fastquadtree
86
87
  ```python
87
88
  from fastquadtree import QuadTree # Point handling
88
89
  from fastquadtree import RectQuadTree # Bounding box handling
90
+ from fastquadtree.pyqtree import Index # Drop-in pyqtree shim (6.567x faster while keeping the same API)
89
91
  ```
90
92
 
91
93
  ## Benchmarks
@@ -30,6 +30,7 @@ Rust-optimized quadtree with a clean Python API
30
30
  - Fast KNN and range queries
31
31
  - Optional object tracking for id ↔ object mapping
32
32
  - [100% test coverage](https://codecov.io/gh/Elan456/fastquadtree) and CI on GitHub Actions
33
+ - Offers a drop-in [pyqtree shim](https://elan456.github.io/fastquadtree/benchmark/#pyqtree-drop-in-shim-performance-gains) that is 6.567x faster while keeping the same API
33
34
 
34
35
  ----
35
36
 
@@ -47,6 +48,7 @@ pip install fastquadtree
47
48
  ```python
48
49
  from fastquadtree import QuadTree # Point handling
49
50
  from fastquadtree import RectQuadTree # Bounding box handling
51
+ from fastquadtree.pyqtree import Index # Drop-in pyqtree shim (6.567x faster while keeping the same API)
50
52
  ```
51
53
 
52
54
  ## Benchmarks
@@ -58,15 +58,15 @@ def bench_shim(points, queries, *, track_objects: bool, with_objs: bool):
58
58
  )
59
59
  if with_objs:
60
60
  for i, p in enumerate(points):
61
- qt.insert(p, id_=i, obj=i) # store a tiny object
61
+ qt.insert(p, obj=i) # store a tiny object
62
62
  else:
63
- for i, p in enumerate(points):
64
- qt.insert(p, id_=i)
63
+ for _, p in enumerate(points):
64
+ qt.insert(p)
65
65
  t_build = now() - t0
66
66
 
67
67
  t0 = now()
68
68
  for q in queries:
69
- _ = qt.query(q) # tuples path for speed
69
+ _ = qt.query(q, as_items=track_objects) # tuples path for speed
70
70
  t_query = now() - t0
71
71
  return t_build, t_query
72
72
 
@@ -188,18 +188,33 @@ def main():
188
188
  | Native | {fmt(n_build)} | {fmt(n_query)} | {fmt(n_build + n_query)} |
189
189
  | Shim (no tracking) | {fmt(s_build_no_map)} | {fmt(s_query_no_map)} | {fmt(s_build_no_map + s_query_no_map)} |
190
190
  | Shim (tracking) | {fmt(s_build_map)} | {fmt(s_query_map)} | {fmt(s_build_map + s_query_map)} |
191
- | pyqtree (fastquadtree) | {fmt(fqt_build)} | {fmt(fqt_query)} | {fmt(fqt_build + fqt_query)} |
192
- | pyqtree (original) | {fmt(p_build)} | {fmt(p_query)} | {fmt(p_build + p_query)} |
193
191
 
194
192
  ### Summary
195
193
 
196
194
  Using the shim with object tracking increases build time by {fmt(s_build_map / n_build)}x and query time by {fmt(s_query_map / n_query)}x.
197
195
  **Total slowdown = {fmt((s_build_map + s_query_map) / (n_build + n_query))}x.**
198
196
 
197
+ Adding the object map tends to only impact the build time, not the query time.
198
+
199
+ ## pyqtree drop-in shim performance gains
200
+
201
+ ### Configuration
202
+ - Points: {args.points:,}
203
+ - Queries: {args.queries}
204
+ - Repeats: {args.repeats}
205
+
206
+ ### Results
207
+
208
+ | Variant | Build | Query | Total |
209
+ |---|---:|---:|---:|
210
+ | pyqtree (fastquadtree) | {fmt(fqt_build)} | {fmt(fqt_query)} | {fmt(fqt_build + fqt_query)} |
211
+ | pyqtree (original) | {fmt(p_build)} | {fmt(p_query)} | {fmt(p_build + p_query)} |
212
+
213
+ ### Summary
214
+
199
215
  If you directly replace pyqtree with the drop-in `fastquadtree.pyqtree.Index` shim, you get a build time of {fmt(fqt_build)}s and query time of {fmt(fqt_query)}s.
200
- This is a total speedup of {fmt((p_build + p_query) / (fqt_build + fqt_query))}x compared to the original pyqtree and requires no code changes.
216
+ This is a **total speedup of {fmt((p_build + p_query) / (fqt_build + fqt_query))}x** compared to the original pyqtree and requires no code changes.
201
217
 
202
- Adding the object map only impacts the build time, not the query time.
203
218
  """
204
219
  print(md.strip())
205
220
 
@@ -0,0 +1,4 @@
1
+ # fastquadtree.pyqtree
2
+ ::: fastquadtree.pyqtree.Index
3
+ options:
4
+ inherited_members: true
@@ -46,16 +46,35 @@ Quadtrees are the focus of the benchmark, but Rtrees are included for reference.
46
46
 
47
47
  | Variant | Build | Query | Total |
48
48
  |---|---:|---:|---:|
49
- | Native | 0.181 | 2.024 | 2.205 |
50
- | Shim (no map) | 0.301 | 1.883 | 2.184 |
51
- | Shim (track+objs) | 0.651 | 2.016 | 2.667 |
49
+ | Native | 0.134 | 1.873 | 2.007 |
50
+ | Shim (no tracking) | 0.245 | 1.815 | 2.060 |
51
+ | Shim (tracking) | 0.714 | 2.062 | 2.776 |
52
52
 
53
53
  ### Summary
54
54
 
55
- Using the shim with object tracking increases build time by 3.604x and query time by 0.996x.
56
- **Total slowdown = 1.210x.**
55
+ Using the shim with object tracking increases build time by 5.345x and query time by 1.101x.
56
+ **Total slowdown = 1.383x.**
57
57
 
58
- Adding the object map only impacts the build time, not the query time.
58
+ Adding the object map tends to impact the build time more than query time.
59
+
60
+ ## pyqtree drop-in shim performance gains
61
+
62
+ ### Configuration
63
+ - Points: 500,000
64
+ - Queries: 500
65
+ - Repeats: 5
66
+
67
+ ### Results
68
+
69
+ | Variant | Build | Query | Total |
70
+ |---|---:|---:|---:|
71
+ | pyqtree (fastquadtree) | 0.443 | 2.007 | 2.450 |
72
+ | pyqtree (original) | 3.233 | 12.857 | 16.091 |
73
+
74
+ ### Summary
75
+
76
+ If you directly replace pyqtree with the drop-in `fastquadtree.pyqtree.Index` shim, you get a build time of 0.443s and query time of 2.007s.
77
+ This is a **total speedup of 6.567x** compared to the original pyqtree and requires no code changes.
59
78
 
60
79
  ---------
61
80
 
@@ -52,8 +52,10 @@ See examples of how fastquadtree can be used in the [runnables](runnables.md) se
52
52
  pip install fastquadtree
53
53
  ```
54
54
 
55
+ ## Import
56
+
55
57
  ```python
56
58
  from fastquadtree import QuadTree # Point handling
57
59
  from fastquadtree import RectQuadTree # Bounding box handling
58
- from fastquadtree.pyqtree import Index # Drop-in replacement for pyqtree
60
+ from fastquadtree.pyqtree import Index # Drop-in replacement for pyqtree (6.567x faster while keeping the same API)
59
61
  ```
@@ -20,10 +20,10 @@ from fastquadtree import QuadTree
20
20
  # 1) Make a tree that covers your world
21
21
  qt = QuadTree(bounds=(0, 0, 1000, 1000), capacity=20)
22
22
 
23
- # 2) Add some stuff
24
- a = qt.insert((10, 10)) # auto id
25
- b = qt.insert((200, 300)) # auto id
26
- _ = qt.insert((999, 500), id_=42) # you can choose ids too
23
+ # 2) Add some stuff (a, b, and c are auto-generated ids)
24
+ a = qt.insert((10, 10))
25
+ b = qt.insert((200, 300))
26
+ c = qt.insert((999, 500))
27
27
 
28
28
  # 3) Ask spatial questions
29
29
  print("Range hits:", qt.query((0, 0, 250, 350))) # -> [(id, x, y), ...]
@@ -100,7 +100,7 @@ Tip: leave `track_objects=False` for max speed when you do not need object mappi
100
100
  Keep the same `QuadTree` instance alive for UIs or game loops. Wipe contents and optionally reset ids.
101
101
 
102
102
  ```python
103
- qt.clear(reset_ids=True) # tree is empty, auto ids start again at 1
103
+ qt.clear() # tree is empty, auto ids start again at 0, all objects forgotten
104
104
  ```
105
105
 
106
106
  ---
@@ -69,10 +69,11 @@ nav:
69
69
  - Runnables: runnables.md
70
70
  - Benchmark: benchmark.md
71
71
  - API:
72
- - fastquadtree.QuadTree: api/quadtree.md
73
- - fastquadtree.RectQuadTree: api/rect_quadtree.md
74
- - fastquadtree.RectItem: api/rect_item.md
75
- - fastquadtree.PointItem: api/point_item.md
72
+ - QuadTree: api/quadtree.md
73
+ - RectQuadTree: api/rect_quadtree.md
74
+ - pyqtree: api/pyqtree.md
75
+ - RectItem: api/rect_item.md
76
+ - PointItem: api/point_item.md
76
77
 
77
78
  extra:
78
79
  social:
@@ -0,0 +1,258 @@
1
+ # _abc_quadtree.py
2
+ from __future__ import annotations
3
+
4
+ from abc import ABC, abstractmethod
5
+ from typing import Any, Generic, Iterable, Tuple, TypeVar
6
+
7
+ from ._item import Item # base class for PointItem and RectItem
8
+ from ._obj_store import ObjStore
9
+
10
+ Bounds = Tuple[float, float, float, float]
11
+
12
+ # Generic parameters
13
+ G = TypeVar("G") # geometry type, e.g. Point or Bounds
14
+ HitT = TypeVar("HitT") # raw native tuple, e.g. (id,x,y) or (id,x0,y0,x1,y1)
15
+ ItemType = TypeVar("ItemType", bound=Item) # e.g. PointItem or RectItem
16
+
17
+
18
+ class _BaseQuadTree(Generic[G, HitT, ItemType], ABC):
19
+ """
20
+ Shared logic for Python QuadTree wrappers over native Rust engines.
21
+
22
+ Concrete subclasses must implement:
23
+ - _new_native(bounds, capacity, max_depth)
24
+ - _make_item(id_, geom, obj)
25
+ """
26
+
27
+ __slots__ = (
28
+ "_bounds",
29
+ "_capacity",
30
+ "_count",
31
+ "_max_depth",
32
+ "_native",
33
+ "_next_id",
34
+ "_store",
35
+ "_track_objects",
36
+ )
37
+
38
+ # ---- required native hooks ----
39
+
40
+ @abstractmethod
41
+ def _new_native(self, bounds: Bounds, capacity: int, max_depth: int | None) -> Any:
42
+ """Create the native engine instance."""
43
+
44
+ @abstractmethod
45
+ def _make_item(self, id_: int, geom: G, obj: Any | None) -> ItemType:
46
+ """Build an ItemType from id, geometry, and optional object."""
47
+
48
+ # ---- ctor ----
49
+
50
+ def __init__(
51
+ self,
52
+ bounds: Bounds,
53
+ capacity: int,
54
+ *,
55
+ max_depth: int | None = None,
56
+ track_objects: bool = False,
57
+ ):
58
+ self._bounds = bounds
59
+ self._max_depth = max_depth
60
+ self._capacity = capacity
61
+ self._native = self._new_native(bounds, capacity, max_depth)
62
+
63
+ self._track_objects = bool(track_objects)
64
+ self._store: ObjStore[ItemType] | None = ObjStore() if track_objects else None
65
+
66
+ # Auto ids when not using ObjStore.free slots
67
+ self._next_id = 0
68
+ self._count = 0
69
+
70
+ # ---- internal helper ----
71
+
72
+ def _ids_to_objects(self, ids: Iterable[int]) -> list[Any]:
73
+ """Map ids -> Python objects via ObjStore in a batched way."""
74
+ if self._store is None:
75
+ raise ValueError("Cannot map ids to objects when track_objects=False")
76
+ return self._store.get_many_objects(list(ids))
77
+
78
+ # ---- shared API ----
79
+
80
+ def insert(self, geom: G, *, obj: Any | None = None) -> int:
81
+ """
82
+ Insert a single item.
83
+
84
+ Args:
85
+ geom: Point (x, y) or Rect (x0, y0, x1, y1) depending on quadtree type.
86
+ obj: Optional Python object to associate with id if tracking is enabled.
87
+
88
+ Returns:
89
+ The id used for this insert.
90
+
91
+ Raises:
92
+ ValueError: If geometry is outside the tree bounds.
93
+ """
94
+ if self._store is not None:
95
+ # Reuse a dense free slot if available, else append
96
+ rid = self._store.alloc_id()
97
+ else:
98
+ rid = self._next_id
99
+ self._next_id += 1
100
+
101
+ if not self._native.insert(rid, geom):
102
+ bx0, by0, bx1, by1 = self._bounds
103
+ raise ValueError(
104
+ f"Geometry {geom!r} is outside bounds ({bx0}, {by0}, {bx1}, {by1})"
105
+ )
106
+
107
+ if self._store is not None:
108
+ self._store.add(self._make_item(rid, geom, obj))
109
+
110
+ self._count += 1
111
+ return rid
112
+
113
+ def insert_many(self, geoms: list[G], objs: list[Any] | None = None) -> int:
114
+ """
115
+ Bulk insert with auto-assigned contiguous ids. Faster than inserting one-by-one.<br>
116
+
117
+ If tracking is enabled, the objects will be bulk stored internally.
118
+ If no objects are provided, the items will have obj=None (if tracking).
119
+
120
+ Args:
121
+ geoms: List of geometries.
122
+ objs: Optional list of Python objects aligned with geoms.
123
+
124
+ Returns:
125
+ Number of items inserted.
126
+
127
+ Raises:
128
+ ValueError: If any geometry is outside bounds.
129
+ """
130
+ if not geoms:
131
+ return 0
132
+
133
+ if self._store is None:
134
+ # Simple contiguous path with native bulk insert
135
+ start_id = self._next_id
136
+ last_id = self._native.insert_many(start_id, geoms)
137
+ num = last_id - start_id + 1
138
+ if num < len(geoms):
139
+ raise ValueError("One or more items are outside tree bounds")
140
+ self._next_id = last_id + 1
141
+ self._count += num
142
+ return num
143
+
144
+ # With tracking enabled:
145
+ start_id = len(self._store._arr) # contiguous tail position
146
+ last_id = self._native.insert_many(start_id, geoms)
147
+ num = last_id - start_id + 1
148
+ if num < len(geoms):
149
+ raise ValueError("One or more items are outside tree bounds")
150
+
151
+ # Add items to the store in one pass
152
+ if objs is None:
153
+ for off, geom in enumerate(geoms):
154
+ id_ = start_id + off
155
+ self._store.add(self._make_item(id_, geom, None))
156
+ else:
157
+ if len(objs) != len(geoms):
158
+ raise ValueError("objs length must match geoms length")
159
+ for off, (geom, o) in enumerate(zip(geoms, objs)):
160
+ id_ = start_id + off
161
+ self._store.add(self._make_item(id_, geom, o))
162
+
163
+ # Keep _next_id monotonic for the non-tracking path
164
+ self._next_id = max(self._next_id, last_id + 1)
165
+
166
+ self._count += num
167
+ return num
168
+
169
+ def delete(self, id_: int, geom: G) -> bool:
170
+ """
171
+ Delete an item by id and exact geometry.
172
+
173
+ Returns:
174
+ True if the item was found and deleted.
175
+ """
176
+ deleted = self._native.delete(id_, geom)
177
+ if deleted:
178
+ self._count -= 1
179
+ if self._store is not None:
180
+ self._store.pop_id(id_)
181
+ return deleted
182
+
183
+ def attach(self, id_: int, obj: Any) -> None:
184
+ """
185
+ Attach or replace the Python object for an existing id.
186
+ Tracking must be enabled.
187
+ """
188
+ if self._store is None:
189
+ raise ValueError("Cannot attach objects when track_objects=False")
190
+ it = self._store.by_id(id_)
191
+ if it is None:
192
+ raise KeyError(f"Id {id_} not found in quadtree")
193
+ # Preserve geometry from existing item
194
+ self._store.add(self._make_item(id_, it.geom, obj)) # type: ignore[attr-defined]
195
+
196
+ def delete_by_object(self, obj: Any) -> bool:
197
+ """
198
+ Delete an item by Python object identity. Tracking must be enabled.
199
+ """
200
+ if self._store is None:
201
+ raise ValueError("Cannot delete by object when track_objects=False")
202
+ it = self._store.by_obj(obj)
203
+ if it is None:
204
+ return False
205
+ return self.delete(it.id_, it.geom) # type: ignore[arg-type]
206
+
207
+ def clear(self) -> None:
208
+ """
209
+ Empty the tree in place, preserving bounds, capacity, and max_depth.
210
+
211
+ If tracking is enabled, the id -> object mapping is also cleared.
212
+ The ids are reset to start at zero again.
213
+ """
214
+ self._native = self._new_native(self._bounds, self._capacity, self._max_depth)
215
+ self._count = 0
216
+ if self._store is not None:
217
+ self._store.clear()
218
+ self._next_id = 0
219
+
220
+ def get_all_objects(self) -> list[Any]:
221
+ """
222
+ Return all tracked Python objects in the tree.
223
+ """
224
+ if self._store is None:
225
+ raise ValueError("Cannot get objects when track_objects=False")
226
+ return [t.obj for t in self._store.items() if t.obj is not None]
227
+
228
+ def get_all_items(self) -> list[ItemType]:
229
+ """
230
+ Return all Item wrappers in the tree.
231
+ """
232
+ if self._store is None:
233
+ raise ValueError("Cannot get items when track_objects=False")
234
+ return list(self._store.items())
235
+
236
+ def get_all_node_boundaries(self) -> list[Bounds]:
237
+ """
238
+ Return all node boundaries in the tree. Useful for visualization.
239
+ """
240
+ return self._native.get_all_node_boundaries()
241
+
242
+ def get(self, id_: int) -> Any | None:
243
+ """
244
+ Return the object associated with id, if tracking is enabled.
245
+ """
246
+ if self._store is None:
247
+ raise ValueError("Cannot get objects when track_objects=False")
248
+ item = self._store.by_id(id_)
249
+ return None if item is None else item.obj
250
+
251
+ def count_items(self) -> int:
252
+ """
253
+ Return the number of items currently in the tree (native count).
254
+ """
255
+ return self._native.count_items()
256
+
257
+ def __len__(self) -> int:
258
+ return self._count
@@ -0,0 +1,167 @@
1
+ # _bimap.py
2
+ from __future__ import annotations
3
+
4
+ from operator import itemgetter
5
+ from typing import Any, Generic, Iterable, Iterator, TypeVar
6
+
7
+ from ._item import Item # base class for PointItem and RectItem
8
+
9
+ TItem = TypeVar("TItem", bound=Item)
10
+
11
+
12
+ class ObjStore(Generic[TItem]):
13
+ """
14
+ High-performance id <-> object store for dense, auto-assigned ids.
15
+
16
+ Storage
17
+ - _arr[id] -> Item or None
18
+ - _objs[id] -> Python object or None
19
+ - _obj_to_id: reverse identity map id(obj) -> id
20
+ - _free: LIFO free-list of reusable ids
21
+
22
+ Assumptions
23
+ - Ids are assigned by the shim and are dense [0..len) with possible holes
24
+ created by deletes. New inserts reuse holes via the free-list.
25
+ """
26
+
27
+ __slots__ = ("_arr", "_free", "_len", "_obj_to_id", "_objs")
28
+
29
+ def __init__(self, items: Iterable[TItem] | None = None) -> None:
30
+ self._arr: list[TItem | None] = []
31
+ self._objs: list[Any | None] = []
32
+ self._obj_to_id: dict[int, int] = {}
33
+ self._free: list[int] = [] # LIFO
34
+ self._len: int = 0 # live items
35
+
36
+ if items:
37
+ for it in items:
38
+ self.add(it)
39
+
40
+ # -------- core --------
41
+
42
+ def add(self, item: TItem) -> None:
43
+ """
44
+ Insert or replace the mapping at item.id_. Reverse map updated so obj points to id.
45
+ """
46
+ id_ = item.id_
47
+ obj = item.obj
48
+
49
+ # ids must be dense and assigned by the caller
50
+ if id_ > len(self._arr):
51
+ raise AssertionError(
52
+ "ObjStore.add received an out-of-order id, use alloc_id() to get the next available id"
53
+ )
54
+
55
+ if id_ == len(self._arr):
56
+ # append
57
+ self._arr.append(item)
58
+ self._objs.append(obj)
59
+ self._len += 1
60
+ else:
61
+ # replace or fill a hole
62
+ old = self._arr[id_]
63
+ if old is None:
64
+ self._len += 1
65
+ elif old.obj is not None:
66
+ self._obj_to_id.pop(id(old.obj), None)
67
+ self._arr[id_] = item
68
+ self._objs[id_] = obj
69
+
70
+ if obj is not None:
71
+ self._obj_to_id[id(obj)] = id_
72
+
73
+ def by_id(self, id_: int) -> TItem | None:
74
+ return self._arr[id_] if 0 <= id_ < len(self._arr) else None
75
+
76
+ def by_obj(self, obj: Any) -> TItem | None:
77
+ id_ = self._obj_to_id.get(id(obj))
78
+ return self.by_id(id_) if id_ is not None else None
79
+
80
+ def pop_id(self, id_: int) -> TItem | None:
81
+ """Remove by id. Dense ids go to the free-list for reuse."""
82
+ if not (0 <= id_ < len(self._arr)):
83
+ return None
84
+ it = self._arr[id_]
85
+ if it is None:
86
+ return None
87
+ self._arr[id_] = None
88
+ self._objs[id_] = None
89
+ if it.obj is not None:
90
+ self._obj_to_id.pop(id(it.obj), None)
91
+ self._free.append(id_)
92
+ self._len -= 1
93
+ return it
94
+
95
+ # -------- allocation --------
96
+
97
+ def alloc_id(self) -> int:
98
+ """
99
+ Get a reusable dense id. Uses free-list else appends at the tail.
100
+ Build your Item with this id then call add(item).
101
+ """
102
+ return self._free.pop() if self._free else len(self._arr)
103
+
104
+ # -------- fast batch gathers --------
105
+
106
+ def get_many_by_ids(self, ids: Iterable[int], *, chunk: int = 2048) -> list[TItem]:
107
+ """
108
+ Batch: return Items for ids, preserving order.
109
+ Uses C-level itemgetter on the dense array in chunks.
110
+ """
111
+ ids_list: list[int] = list(ids)
112
+ if not ids_list:
113
+ return []
114
+
115
+ out: list[TItem] = []
116
+ extend = out.extend
117
+ arr = self._arr
118
+ for i in range(0, len(ids_list), chunk):
119
+ block = ids_list[i : i + chunk]
120
+ vals = itemgetter(*block)(arr) # tuple or single item
121
+ extend(vals if isinstance(vals, tuple) else (vals,))
122
+ return out
123
+
124
+ def get_many_objects(self, ids: Iterable[int], *, chunk: int = 2048) -> list[Any]:
125
+ """
126
+ Batch: return Python objects for ids, preserving order.
127
+ Mirrors get_many_by_ids but reads from _objs.
128
+ """
129
+ ids_list: list[int] = list(ids)
130
+ if not ids_list:
131
+ return []
132
+
133
+ out: list[Any] = []
134
+ extend = out.extend
135
+ objs = self._objs
136
+ for i in range(0, len(ids_list), chunk):
137
+ block = ids_list[i : i + chunk]
138
+ vals = itemgetter(*block)(objs) # tuple or single object
139
+ extend(vals if isinstance(vals, tuple) else (vals,))
140
+ return out
141
+
142
+ # -------- convenience and iteration --------
143
+
144
+ def __len__(self) -> int:
145
+ return self._len
146
+
147
+ def clear(self) -> None:
148
+ self._arr.clear()
149
+ self._objs.clear()
150
+ self._obj_to_id.clear()
151
+ self._free.clear()
152
+ self._len = 0
153
+
154
+ def contains_id(self, id_: int) -> bool:
155
+ return 0 <= id_ < len(self._arr) and self._arr[id_] is not None
156
+
157
+ def contains_obj(self, obj: Any) -> bool:
158
+ return id(obj) in self._obj_to_id
159
+
160
+ def items_by_id(self) -> Iterator[tuple[int, TItem]]:
161
+ for i, it in enumerate(self._arr):
162
+ if it is not None:
163
+ yield i, it
164
+
165
+ def items(self) -> Iterator[TItem]:
166
+ for _, it in self.items_by_id():
167
+ yield it