diff --git a/src/python_bindings.rs b/src/python_bindings.rs index 07caedb..a67c01b 100644 --- a/src/python_bindings.rs +++ b/src/python_bindings.rs @@ -599,10 +599,8 @@ fn is_tuple_or_list_of_two(obj: &Bound) -> bool { #[cfg(feature = "opencv")] #[pyclass] pub struct ResizeIterator { - /// Raw buffer pointers and their dimensions - buffers: Vec<(Vec, (usize, usize, usize))>, // (buffer, (height, width, channels)) - /// Current iteration index - index: usize, + /// Completed buffers waiting to be transferred to NumPy. + buffers: std::collections::VecDeque<(Vec, (usize, usize, usize))>, } #[cfg(feature = "opencv")] @@ -613,17 +611,13 @@ impl ResizeIterator { } fn __next__<'py>(&mut self, py: Python<'py>) -> Option>> { - if self.index >= self.buffers.len() { - return None; - } - - let (buffer, (height, width, channels)) = &self.buffers[self.index]; - self.index += 1; + let (buffer, (height, width, channels)) = self.buffers.pop_front()?; - // Convert raw buffer directly to PyArray3 - ZERO intermediate steps! - match ndarray::Array3::from_shape_vec((*height, *width, *channels), buffer.clone()) { + // Transfer the completed allocation directly into the ndarray instead + // of cloning the full resized image during iteration. + match ndarray::Array3::from_shape_vec((height, width, channels), buffer) { Ok(array) => Some(PyArray3::from_owned_array_bound(py, array)), - Err(_) => None, // Skip malformed arrays + Err(_) => None, // Malformed buffer; end iteration (StopIteration) } } @@ -677,8 +671,7 @@ pub fn batch_resize_images_iterator<'py>( return Bound::new( py, ResizeIterator { - buffers: Vec::new(), - index: 0, + buffers: std::collections::VecDeque::new(), }, ); } @@ -856,7 +849,12 @@ pub fn batch_resize_images_iterator<'py>( .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("Resize failed: {}", e)))?; // Return iterator with raw buffers - conversion happens on-demand! - Bound::new(py, ResizeIterator { buffers, index: 0 }) + Bound::new( + py, + ResizeIterator { + buffers: buffers.into(), + }, + ) } // TSR CROPPING OPERATIONS (BENCHMARK WINNERS) diff --git a/tests/test_python_bindings.py b/tests/test_python_bindings.py index d00637f..7947306 100644 --- a/tests/test_python_bindings.py +++ b/tests/test_python_bindings.py @@ -156,6 +156,31 @@ def test_zero_copy_resize_iterator_rejects_non_contiguous_array(self): with pytest.raises(ValueError, match="C-contiguous"): tsr.batch_resize_images_iterator([strided], [(16, 16)]) + def test_zero_copy_resize_iterator_transfers_each_result(self): + """Iterator should yield owned results once and report remaining length.""" + if not hasattr(tsr, "batch_resize_images_iterator"): + pytest.skip("OpenCV resize iterator bindings not available") + + images = [ + np.full((32, 48, 3), 17, dtype=np.uint8), + np.full((24, 40, 3), 231, dtype=np.uint8), + ] + iterator = tsr.batch_resize_images_iterator(images, [(12, 8), (10, 6)]) + + assert len(iterator) == 2 + first = next(iterator) + assert first.shape == (8, 12, 3) + assert np.all(first == 17) + assert len(iterator) == 1 + + second = next(iterator) + assert second.shape == (6, 10, 3) + assert np.all(second == 231) + assert len(iterator) == 0 + + with pytest.raises(StopIteration): + next(iterator) + class TestErrorHandling: """Test error handling and edge cases."""