Update PyO3 and minor chores

python/Cargo.toml

@@ -17,9 +17,9 @@ crate-type = ["cdylib"]
 
 [dependencies]
 bincode = { version = "2.0", features = ["serde"] }
-pyo3 = { version = "0.26", features = ["experimental-inspect"] }
+pyo3 = { version = "0.27", features = ["experimental-inspect"] }
 zelll = { path  = "..", features = ["serde"] }
 serde = { version = "1.0" }
 
 [dev-dependencies]
-pyo3-introspection = "0.26.0"
+pyo3-introspection = "0.27.0"

python/examples/psssh.py

@@ -121,6 +121,7 @@ def harmonic_potential(x, r, k):
         compiled = nutpie.compile_pymc_model(model)
         # zelll's Python bindings don't benefit from nutpie's multithreading
         # unless a free-threaded Python build is used
+        # however, num_dual currently has no pre-built free-threaded wheels
         trace = nutpie.sample(
             compiled,
             save_warmup=False,

python/examples/requirements.txt

@@ -4,3 +4,4 @@ numpy==2.2.6
 nutpie==0.15.2
 pymc==5.25.1
 pytensor==2.31.7
+numba==0.63.1

python/src/lib.rs

@@ -12,11 +12,11 @@ use pyo3::types::PyIterator;
 // TODO: where we accept an arbitrary `ParticlesIterable` as we currently do
 // TODO: but also perhaps clone into some buffer so we don't have to hold the GIL for too long
 #[derive(Clone)]
-struct ParticlesIterable<'py> {
-    inner: Bound<'py, PyAny>,
+struct ParticlesIterable<'a, 'py> {
+    inner: Borrowed<'a, 'py, PyAny>,
 }
 
-impl<'py> IntoIterator for ParticlesIterable<'py> {
+impl<'a, 'py> IntoIterator for ParticlesIterable<'a, 'py> {
     type Item = [f64; 3];
     type IntoIter = ParticlesIterator<'py>;
 
@@ -45,7 +45,7 @@ impl<'py> Iterator for ParticlesIterator<'py> {
         // otherwise, retry with next element
         loop {
             match self.iter.next().transpose() {
-                Ok(Some(p)) => match <[f64; 3] as FromPyObject>::extract_bound(&p) {
+                Ok(Some(p)) => match <[f64; 3] as FromPyObject>::extract(p.as_borrowed()) {
                     Ok(p) => break Some(p),
                     Err(_) => (),
                 },
@@ -109,12 +109,14 @@ impl PyCellGrid {
     /// see `CellGrid`.
     #[new]
     #[pyo3(signature = (particles: "typing.Iterable[typing.Sequence[float]] | None" = None, /, cutoff=1.0) -> "zelll.CellGrid")]
-    fn new<'py>(py: Python<'py>, particles: Option<&Bound<'py, PyAny>>, cutoff: f64) -> Self {
+    fn new<'py>(py: Python<'py>, particles: Option<Bound<PyAny>>, cutoff: f64) -> Self {
         let inner = match particles {
             Some(p) => {
                 // TODO: see if we can simplify ParticlesIterable, it wraps Bound<>, so holds the GIL
                 // TODO: would like to call ::new() detached from the GIL if that's possible?
-                let particles = ParticlesIterable { inner: p.clone() };
+                let particles = ParticlesIterable {
+                    inner: p.as_borrowed(),
+                };
                 CellGrid::new(particles, cutoff)
             }
             // nutpie+multithreading needs to serialize/deserialize PyCellGrid
@@ -152,11 +154,11 @@ impl PyCellGrid {
     #[pyo3(signature = (particles: "typing.Iterable[typing.Sequence[float]]", /, cutoff=None))]
     fn rebuild<'py>(
         mut slf: PyRefMut<'_, Self>,
-        particles: &Bound<'py, PyAny>,
+        particles: Bound<'py, PyAny>,
         cutoff: Option<f64>,
     ) -> () {
         let particles = ParticlesIterable {
-            inner: particles.clone(),
+            inner: particles.as_borrowed(),
         };
 
         slf.inner.rebuild_mut(particles, cutoff);
@@ -201,9 +203,9 @@ impl PyCellGrid {
     #[pyo3(signature = (coordinates: "typing.Sequence[float]") -> "zelll.CellQueryIter | None")]
     fn query_neighbors(
         slf: PyRef<'_, Self>,
-        coordinates: &Bound<'_, PyAny>,
+        coordinates: Bound<'_, PyAny>,
     ) -> Option<PyCellQueryIter> {
-        PyCellQueryIter::new(slf, coordinates)
+        PyCellQueryIter::new(slf, coordinates.as_borrowed())
     }
 
     /// Given 3D `coordinates`, this returns a list over all particles in the neighborhood
@@ -270,9 +272,7 @@ impl PyCellGrid {
 #[pyclass(name = "CellGridIter", module = "zelll", unsendable)]
 pub struct PyCellGridIter {
     // TODO: it looks like we probably don't need `_owner`
-    // TODO: also haven't yet figured out how to migrate this from IntoPy to IntoPyObject
-    // TODO: likely need more unsafe using Py::from_*_pointer()
-    // _owner: PyObject,
+    // _owner: Py<PyAny>,
     // TODO: `_keep_borrow` is enough to maintain correct drop order *and* prevents `PyCellGrid`
     // TODO: from being mutated while `PyCellGridIter` is still alive
     _keep_borrow: PyRef<'static, PyCellGrid>,
@@ -282,9 +282,11 @@ pub struct PyCellGridIter {
 impl PyCellGridIter {
     fn new(py_cellgrid: PyRef<'_, PyCellGrid>) -> Self {
         // let py = py_cellgrid.py();
-        // let _owner = (&py_cellgrid).into_py(py);
+        // let _owner = (&py_cellgrid)
+        //     .into_py_any(py)
+        //     .expect("could not store owner internally");
         let iter = Box::new((&py_cellgrid).inner.particle_pairs());
-        // SAFETY: but the idea is that `_keep_borrow` makes sure that `iter`s lifetime can be extended
+        // SAFETY: the idea is that `_keep_borrow` makes sure that `iter`s lifetime can be extended
         // SAFETY: replicating some ideas from
         // SAFETY: https://github.com/PyO3/pyo3/issues/1085 and
         // SAFETY: https://github.com/PyO3/pyo3/issues/1089
@@ -324,6 +326,12 @@ impl PyCellGridIter {
 //     }
 // }
 
+// impl Drop for PyCellQueryIter {
+//     fn drop(&mut self) {
+//         eprintln!("Dropping PyCellQueryIter");
+//     }
+// }
+
 #[pymethods]
 impl PyCellGridIter {
     fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
@@ -352,8 +360,11 @@ pub struct PyCellQueryIter {
 }
 
 impl PyCellQueryIter {
-    fn new(py_cellgrid: PyRef<'_, PyCellGrid>, coordinates: &Bound<'_, PyAny>) -> Option<Self> {
-        let coordinates = <[f64; 3] as FromPyObject>::extract_bound(coordinates).ok()?;
+    fn new(
+        py_cellgrid: PyRef<'_, PyCellGrid>,
+        coordinates: Borrowed<'_, '_, PyAny>,
+    ) -> Option<Self> {
+        let coordinates = <[f64; 3] as FromPyObject>::extract(coordinates).ok()?;
         let iter = Box::new((&py_cellgrid).inner.query_neighbors(coordinates)?);
         // SAFETY: see PyCellGridIter
         let iter = unsafe {