feat: require numpy arrays in insert many to use the same dtype as the quadtree

Author: Elan456

Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name = "fastquadtree"
-version = "1.3.0"
+version = "1.3.1"
 edition = "2021"
 
 [lib]

Cargo.toml

@@ -29,6 +29,14 @@
 HitT = TypeVar("HitT")  # raw native tuple, e.g. (id,x,y) or (id,x0,y0,x1,y1)
 ItemType = TypeVar("ItemType", bound=Item)  # e.g. PointItem or RectItem
 
+# Quadtree dtype to numpy dtype mapping
+QUADTREE_DTYPE_TO_NP_DTYPE = {
+    "f32": "float32",
+    "f64": "float64",
+    "i32": "int32",
+    "i64": "int64",
+}
+
 
 def _is_np_array(x: Any) -> bool:
     mod = getattr(x.__class__, "__module__", "")
@@ -250,7 +258,7 @@ def insert_many(
     ) -> int:
         """
         Bulk insert with auto-assigned contiguous ids. Faster than inserting one-by-one.<br>
-        Can accept either a Python sequence of geometries or a NumPy array of shape (N,2) or (N,4) with dtype float32.
+        Can accept either a Python sequence of geometries or a NumPy array of shape (N,2) or (N,4) with a dtype that matches the quadtree's dtype.
 
         If tracking is enabled, the objects will be bulk stored internally.
         If no objects are provided, the items will have obj=None (if tracking).
@@ -289,8 +297,12 @@ def insert_many(
             if geoms.size == 0:
                 return 0
 
-            if geoms.dtype != _np.float32:
-                raise TypeError("Numpy array must use dtype float32")
+            # Check if dtype matches quadtree dtype
+            expected_np_dtype = QUADTREE_DTYPE_TO_NP_DTYPE.get(self._dtype)
+            if geoms.dtype != expected_np_dtype:
+                raise TypeError(
+                    f"Numpy array dtype {geoms.dtype} does not match quadtree dtype {self._dtype}"
+                )
 
         if self._store is None:
             # Simple contiguous path with native bulk insert

pysrc/fastquadtree/_base_quadtree.py

@@ -161,15 +161,15 @@ def _new_native(self, bounds: Bounds, capacity: int, max_depth: int | None) -> A
         """Create the native engine instance."""
         rust_cls = DTYPE_MAP.get(self._dtype)
         if rust_cls is None:
-            raise ValueError(f"Unsupported dtype: {self._dtype}")
+            raise TypeError(f"Unsupported dtype: {self._dtype}")
         return rust_cls(bounds, capacity, max_depth)
 
     @classmethod
     def _new_native_from_bytes(cls, data: bytes, dtype: str = "f32") -> Any:
         """Create a new native engine instance from serialized bytes."""
         rust_cls = DTYPE_MAP.get(dtype)
         if rust_cls is None:
-            raise ValueError(f"Unsupported dtype: {dtype}")
+            raise TypeError(f"Unsupported dtype: {dtype}")
         return rust_cls.from_bytes(data)
 
     @staticmethod

pysrc/fastquadtree/point_quadtree.py

@@ -101,15 +101,15 @@ def _new_native(self, bounds: Bounds, capacity: int, max_depth: int | None) -> A
         """Create the native engine instance."""
         rust_cls = DTYPE_MAP.get(self._dtype)
         if rust_cls is None:
-            raise ValueError(f"Unsupported dtype: {self._dtype}")
+            raise TypeError(f"Unsupported dtype: {self._dtype}")
         return rust_cls(bounds, capacity, max_depth)
 
     @classmethod
     def _new_native_from_bytes(cls, data: bytes, dtype: str = "f32") -> Any:
         """Create a new native engine instance from serialized bytes."""
         rust_cls = DTYPE_MAP.get(dtype)
         if rust_cls is None:
-            raise ValueError(f"Unsupported dtype: {dtype}")
+            raise TypeError(f"Unsupported dtype: {dtype}")
         return rust_cls.from_bytes(data)
 
     @staticmethod

pysrc/fastquadtree/rect_quadtree.py

@@ -52,6 +52,48 @@ def test_type_error_on_wrong_dtype():
     assert len(qt) == 0
 
 
+def test_non_default_dtype_insert_many():
+    qt = QuadTree(BOUNDS, capacity=8, track_objects=True, dtype="f64")
+    points = np.array([[10, 10], [20, 20], [30, 30]], dtype=np.float64)
+    n = qt.insert_many(points)
+    assert n == 3
+    assert len(qt) == 3
+
+    raw = qt.query((0, 0, 40, 40), as_items=False)
+
+    assert len(raw) == 3
+    # ids and positions match
+    m_raw = {t[0]: (t[1], t[2]) for t in raw}
+    for t in raw:
+        assert (t[1], t[2]) == m_raw[t[0]]
+
+
+def test_non_default_quadtree_dtype_with_default_numpy_dtype_raises():
+    qt = QuadTree(BOUNDS, capacity=8, track_objects=True, dtype="f64")
+    points = np.array([[10, 10], [20, 20], [30, 30]], dtype=np.float32)  # Wrong dtype
+    with pytest.raises(TypeError):
+        qt.insert_many(points)
+    assert len(qt) == 0
+
+
+def test_unspported_quadtree_dtype_insert_many_raises():
+    qt = QuadTree(BOUNDS, capacity=8, track_objects=True, dtype="i32")
+    points = np.array([[10, 10], [20, 20], [30, 30]], dtype=np.float32)  # Wrong dtype
+    with pytest.raises(TypeError):
+        qt.insert_many(points)
+    assert len(qt) == 0
+
+    points = np.array(
+        [[10, 10], [20, 20], [30, 30]], dtype=np.uint32
+    )  # unsupported dtype
+    with pytest.raises(TypeError):
+        qt.insert_many(points)
+
+    # QT is also unsupported
+    with pytest.raises(TypeError):
+        qt = QuadTree(BOUNDS, capacity=8, track_objects=True, dtype="u32")
+
+
 def test_insert_empty_numpy_array():
     qt = QuadTree(BOUNDS, capacity=8, track_objects=True)
     points = np.empty((0, 2), dtype=np.float32)

tests/test_insert_many_numpy.py

@@ -5,13 +5,13 @@
 
 def test_unsupported_dtype():
     """Test that providing an unsupported dtype raises ValueError."""
-    with pytest.raises(ValueError):
+    with pytest.raises(TypeError):
         QuadTree((0, 0, 100, 100), capacity=4, track_objects=True, dtype="f128")  # type: ignore
 
     # From bytes
     qt = QuadTree((0, 0, 100, 100), capacity=4, track_objects=True, dtype="f32")
     data = qt.to_bytes()
-    with pytest.raises(ValueError):
+    with pytest.raises(TypeError):
         QuadTree.from_bytes(data, dtype="f128")  # type: ignore
 
 

tests/test_point_quadtree_dtypes.py

@@ -195,7 +195,7 @@ def test_accurate_obj_output_with_tracking():
 
 def test_unsupported_dtype():
     """Test that providing an unsupported dtype raises ValueError."""
-    with pytest.raises(ValueError):
+    with pytest.raises(TypeError):
         rq.RectQuadTree(
             b_to_float(0, 0, 100, 100), capacity=4, track_objects=True, dtype="f128"
         )  # type: ignore
@@ -205,7 +205,7 @@ def test_unsupported_dtype():
         b_to_float(0, 0, 100, 100), capacity=4, track_objects=True, dtype="f32"
     )
     data = qt.to_bytes()
-    with pytest.raises(ValueError):
+    with pytest.raises(TypeError):
         rq.RectQuadTree.from_bytes(data, dtype="f128")  # type: ignore