src/python.rs

changeset 3
fbdee8e4a78d
parent 0
e8f3b6c55ce7
child 6
aefacd832408
--- a/src/python.rs	Fri Nov 28 14:46:35 2025 -0500
+++ b/src/python.rs	Thu Dec 04 14:36:22 2025 -0500
@@ -50,7 +50,20 @@
             }*/
 
             fn __iter__(slf: Bound<'_, Self>) -> PyResult<Py<$iter>> {
-                Py::new(slf.py(), $iter { measure_ref: slf.unbind(), next: 0 })
+                Py::new(slf.py(), $iter {
+                    measure_ref: slf.unbind(),
+                    next: 0,
+                    pad: false,
+                })
+            }
+
+            /// Same as iterating the object, but pads (or cuts) location vectors to 3 dimensions.
+            fn iter_padded(slf: Bound<'_, Self>) -> PyResult<Py<$iter>> {
+                Py::new(slf.py(), $iter {
+                    measure_ref: slf.unbind(),
+                    next: 0,
+                    pad: true,
+                })
             }
         }
 
@@ -61,6 +74,7 @@
         pub struct $iter {
             measure_ref: Py<$name>,
             next: usize,
+            pad: bool,
         }
 
         #[pymethods]
@@ -77,11 +91,22 @@
                 let py = slf.py();
                 let meas_: PyRef<'_, $name> = slf.measure_ref.extract(py)?;
                 let meas = &(meas_.inner);
+                let pad = slf.pad;
                 let next = &mut slf.next;
                 Ok((*next < meas.len()).then(|| {
                     let δ = meas[*next];
                     *next += 1;
-                    (PyArray::from_slice(py, &δ.x.0), δ.α)
+                    if pad {
+                        (
+                            PyArray::from_iter(
+                                py,
+                                δ.x.iter().copied().chain(std::iter::repeat(0.0)).take(3),
+                            ),
+                            δ.α,
+                        )
+                    } else {
+                        (PyArray::from_slice(py, &δ.x.0), δ.α)
+                    }
                 }))
             }
         }

mercurial