Add a context manager for AvailableData

src/runtime.rs

@@ -287,6 +287,35 @@ impl AvailableData {
     fn __iter__(slf: PyRef<'_, Self>) -> PyResult<Py<VideoFrameIterator>> {
         Py::new(slf.py(), slf.frames())
     }
+
+    fn __enter__(&self) -> PyResult<AvailableData> {
+        Ok(AvailableData {
+            inner: self.inner.clone(),
+        })
+    }
+
+    fn __exit__(&self, _exc_type: Option<&PyAny>, _exc_value: Option<&PyAny>, _traceback: Option<&PyAny>) -> bool {
+        let consumed_bytes = self.inner.consumed_bytes.unwrap_or(unsafe {
+            byte_offset_from(self.inner.beg.as_ptr(), self.inner.end.as_ptr())
+        } as usize);
+        log::debug!(
+            "[stream {}] DROP read region: {:p}-{:p}:{}",
+            self.inner.stream_id,
+            self.inner.beg.as_ptr(),
+            self.inner.end.as_ptr(),
+            consumed_bytes
+        );
+        unsafe {
+            capi::acquire_unmap_read(
+                self.inner.runtime.inner.as_ptr(),
+                self.inner.stream_id,
+                consumed_bytes as _,
+            )
+            .ok()
+            .expect("Unexpected failure: Was the CoreRuntime NULL?");
+        }
+        true
+    }
 }
 
 #[pyclass]

tests/test_basic.py

@@ -16,6 +16,8 @@
 from skimage.transform import downscale_local_mean
 import numpy as np
 
+logging.getLogger().setLevel(logging.DEBUG)
+
 
 @pytest.fixture(scope="module")
 def runtime():
@@ -87,28 +89,61 @@ def test_zero_conf_start(runtime: Runtime):
         runtime.start()
 
 
+def test_get_available_data(runtime: Runtime):
+    p = acquire.setup(runtime, "simulated.*empty.*", "Trash")
+    assert p.video[0].camera.identifier is not None
+    assert p.video[0].storage.identifier is not None
+    assert p.video[0].storage.settings.filename == "out.tif"
+    p.video[0].camera.settings.shape = (192, 108)
+    p.video[0].camera.settings.exposure_time_us = 1e4
+    p.video[0].max_frame_count = 10
+    p = runtime.set_configuration(p)
+
+    runtime.start()
+
+    while True:
+        if a := runtime.get_available_data(0):
+            with a:
+                frame_count = a.get_frame_count()
+                assert frame_count > 0
+                logging.info(f"Got {frame_count}")
+
+                i = 0
+                for _ in a.frames():
+                    i += 1
+
+                assert i == frame_count
+            break
+
+    runtime.stop()
+    assert runtime.get_available_data(0) is None
+
+    assert a.get_frame_count() == 0
+
+
 def test_repeat_acq(runtime: Runtime):
     p = acquire.setup(runtime, "simulated: radial sin", "Trash")
     assert p.video[0].camera.identifier is not None
     assert p.video[0].storage.identifier is not None
     assert p.video[0].storage.settings.filename == "out.tif"
     p.video[0].camera.settings.shape = (192, 108)
+    p.video[0].camera.settings.exposure_time_us = 1e4
     p.video[0].max_frame_count = 10
     p = runtime.set_configuration(p)
     runtime.start()
     while True:
         if a := runtime.get_available_data(0):
-            logging.info(f"Got {a.get_frame_count()}")
-            a = None
+            with a:
+                logging.info(f"Got {a.get_frame_count()}")
             break
     runtime.stop()
     assert runtime.get_available_data(0) is None
     # TODO: (nclack) assert 1 acquired frame. stop should block
     runtime.start()
     while True:
         if a := runtime.get_available_data(0):
-            logging.info(f"Got {a.get_frame_count()}")
-            a = None
+            with a:
+                logging.info(f"Got {a.get_frame_count()}")
             break
     runtime.stop()
     assert runtime.get_available_data(0) is None