Context manager for AvailableData

Cargo.toml

@@ -1,32 +1,32 @@
 [package]
 name = "acquire-imaging"
 authors = ["Nathan Clack <[email protected]>"]
-version = "0.2.1"
+version = "0.3.0"
 edition = "2021"
 
 [lib]
 name = "acquire"
 crate-type = ["cdylib"]
 
 [dependencies]
-pyo3 = { version = "0.19", features = [
+pyo3 = { version = "0.20", features = [
     "extension-module",
     "anyhow",
     "abi3-py38",
     "serde",
 ] }
-pyo3-log = "0.8"
-numpy = "0.19"
+pyo3-log = "0.9"
+numpy = "0.20"
 log = "0.4"
 anyhow = "1.0"
 parking_lot = "0.12"
 serde = { version = "1.0", features = ["derive"] }
-pythonize = "0.19"
+pythonize = "0.20.0"
 
 [build-dependencies]
-bindgen = "0.66"
+bindgen = "0.69.1"
 cmake = "0.1"
-http = "0.2"
+http = "1.0"
 json = "0.12"
 reqwest = { version = "0.11", features = ["blocking", "json"] }
 serde = { version = "1.0", features = ["derive"] }

build.rs

@@ -72,7 +72,7 @@ fn main() {
     let bindings = bindgen::Builder::default()
         .header("wrapper.h")
         .clang_arg(format!("-I{}/include", dst.display()))
-        .parse_callbacks(Box::new(bindgen::CargoCallbacks))
+        .parse_callbacks(Box::new(bindgen::CargoCallbacks::new()))
         .generate()
         .expect("Unable to generate bindings");
 

python/acquire/__init__.py

@@ -189,17 +189,18 @@ def is_not_done() -> bool:
         def next_frame() -> Optional[npt.NDArray[Any]]:
             """Get the next frame from the current stream."""
             if nframes[stream_id] < p.video[stream_id].max_frame_count:
-                if packet := runtime.get_available_data(stream_id):
-                    n = packet.get_frame_count()
-                    nframes[stream_id] += n
-                    logging.info(
-                        f"[stream {stream_id}] frame count: {nframes}"
-                    )
-                    f = next(packet.frames())
-                    logging.debug(
-                        f"stream {stream_id} frame {f.metadata().frame_id}"
-                    )
-                    return f.data().squeeze().copy()
+                with runtime.get_available_data(stream_id) as packet:
+                    if packet:
+                        n = packet.get_frame_count()
+                        nframes[stream_id] += n
+                        logging.info(
+                            f"[stream {stream_id}] frame count: {nframes}"
+                        )
+                        f = next(packet.frames())
+                        logging.debug(
+                            f"stream {stream_id} frame {f.metadata().frame_id}"
+                        )
+                        return f.data().squeeze().copy()
             return None
 
         while is_not_done():  # runtime.get_state()==DeviceState.Running:

python/acquire/acquire.pyi

@@ -12,10 +12,18 @@ from typing import (
 
 from numpy.typing import NDArray
 
+@final
+class AvailableDataContext:
+    def __enter__(self) -> Optional[AvailableData]: ...
+    def __exit__(
+        self, exc_type: Any, exc_value: Any, traceback: Any
+    ) -> None: ...
+
 @final
 class AvailableData:
     def frames(self) -> Iterator[VideoFrame]: ...
     def get_frame_count(self) -> int: ...
+    def invalidate(self) -> None: ...
     def __iter__(self) -> Iterator[VideoFrame]: ...
 
 @final
@@ -218,7 +226,7 @@ class Properties:
 class Runtime:
     def __init__(self, *args: None, **kwargs: Any) -> None: ...
     def device_manager(self) -> DeviceManager: ...
-    def get_available_data(self, stream_id: int) -> AvailableData: ...
+    def get_available_data(self, stream_id: int) -> AvailableDataContext: ...
     def get_configuration(self) -> Properties: ...
     def get_capabilities(self) -> Capabilities: ...
     def get_state(self) -> DeviceState: ...

src/lib.rs

@@ -13,7 +13,7 @@ use device_manager::DeviceManager;
 use pyo3::prelude::*;
 use std::ffi::CStr;
 
-use crate::runtime::{AvailableData, VideoFrame, VideoFrameMetadata, VideoFrameTimestamps};
+use crate::runtime::{AvailableData, AvailableDataContext, VideoFrame, VideoFrameMetadata, VideoFrameTimestamps};
 
 trait Status: Copy + Sized {
     fn is_ok(&self) -> bool;
@@ -60,6 +60,7 @@ fn acquire(py: Python, m: &PyModule) -> PyResult<()> {
     m.add_class::<runtime::Runtime>()?;
     m.add_class::<DeviceManager>()?;
     m.add_class::<AvailableData>()?;
+    m.add_class::<AvailableDataContext>()?;
     m.add_class::<VideoFrame>()?;
     m.add_class::<VideoFrameMetadata>()?;
     m.add_class::<VideoFrameTimestamps>()?;

src/runtime.rs

@@ -78,6 +78,22 @@ impl RawRuntime {
         unsafe { capi::acquire_abort(self.inner.as_ptr()) }.ok()?;
         Ok(())
     }
+
+    fn map_read(&self, stream_id: u32) -> Result<(*mut capi::VideoFrame, *mut capi::VideoFrame)> {
+        let mut beg = null_mut();
+        let mut end = null_mut();
+        unsafe {
+            capi::acquire_map_read(self.inner.as_ptr(), stream_id, &mut beg, &mut end).ok()?;
+        }
+        Ok((beg, end))
+    }
+
+    fn unmap_read(&self, stream_id: u32, consumed_bytes: usize) -> Result<()> {
+        unsafe {
+            capi::acquire_unmap_read(self.inner.as_ptr(), stream_id, consumed_bytes).ok()?;
+        }
+        Ok(())
+    }
 }
 
 impl Drop for RawRuntime {
@@ -175,34 +191,11 @@ impl Runtime {
         Python::allow_threads(py, || Ok(self.inner.execute_trigger(stream_id)?))
     }
 
-    fn get_available_data(&self, stream_id: u32) -> PyResult<Option<AvailableData>> {
-        let mut beg = null_mut();
-        let mut end = null_mut();
-        unsafe {
-            capi::acquire_map_read(self.as_ref().as_ptr(), stream_id, &mut beg, &mut end).ok()?;
-        }
-        let nbytes = unsafe { byte_offset_from(beg, end) };
-        if nbytes > 0 {
-            log::trace!(
-                "[stream {}] ACQUIRED {:p}-{:p}:{} bytes",
-                stream_id,
-                beg,
-                end,
-                nbytes
-            )
-        };
-        Ok(if nbytes > 0 {
-            Some(AvailableData {
-                inner: Arc::new(RawAvailableData {
-                    runtime: self.inner.clone(),
-                    beg: NonNull::new(beg).ok_or(anyhow!("Expected non-null buffer"))?,
-                    end: NonNull::new(end).ok_or(anyhow!("Expected non-null buffer"))?,
-                    stream_id,
-                    consumed_bytes: None,
-                }),
-            })
-        } else {
-            None
+    fn get_available_data(&self, stream_id: u32) -> PyResult<AvailableDataContext> {
+        Ok(AvailableDataContext {
+            inner: self.inner.clone(),
+            stream_id,
+            available_data: None,
         })
     }
 }
@@ -271,62 +264,116 @@ impl Drop for RawAvailableData {
             self.end.as_ptr(),
             consumed_bytes
         );
-        unsafe {
-            capi::acquire_unmap_read(
-                self.runtime.inner.as_ptr(),
-                self.stream_id,
-                consumed_bytes as _,
-            )
-            .ok()
-            .expect("Unexpected failure: Was the CoreRuntime NULL?");
-        }
+        self.runtime
+            .unmap_read(self.stream_id, consumed_bytes)
+            .expect("Unexpected failure: Was the runtime NULL?");
     }
 }
 
 #[pyclass]
 pub(crate) struct AvailableData {
-    inner: Arc<RawAvailableData>,
+    inner: Option<Arc<RawAvailableData>>,
 }
 
 #[pymethods]
 impl AvailableData {
     fn get_frame_count(&self) -> usize {
-        self.inner.get_frame_count()
+        if let Some(inner) = &self.inner {
+            inner.get_frame_count()
+        } else {
+            0
+        }
     }
 
     fn frames(&self) -> VideoFrameIterator {
         VideoFrameIterator {
-            store: self.inner.clone(),
-            cur: Mutex::new(self.inner.beg),
-            end: self.inner.end,
+            inner: if let Some(frames) = &self.inner {
+                Some(VideoFrameIteratorInner {
+                    store: frames.clone(),
+                    cur: Mutex::new(frames.beg),
+                    end: frames.end,
+                })
+            } else {
+                None
+            },
         }
     }
 
     fn __iter__(slf: PyRef<'_, Self>) -> PyResult<Py<VideoFrameIterator>> {
         Py::new(slf.py(), slf.frames())
     }
+
+    fn invalidate(&mut self) {
+        // Will drop the RawAvailableData and cause Available data to act like
+        // an empty iterator.
+        self.inner = None;
+    }
 }
 
 #[pyclass]
-struct VideoFrameIterator {
-    store: Arc<RawAvailableData>,
-    cur: Mutex<NonNull<capi::VideoFrame>>,
-    end: NonNull<capi::VideoFrame>,
+pub(crate) struct AvailableDataContext {
+    inner: Arc<RawRuntime>,
+    stream_id: u32,
+    available_data: Option<Py<AvailableData>>,
 }
 
-unsafe impl Send for VideoFrameIterator {}
-
 #[pymethods]
-impl VideoFrameIterator {
-    fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
-        slf
+impl AvailableDataContext {
+    fn __enter__(&mut self) -> PyResult<Option<Py<AvailableData>>> {
+        let AvailableDataContext {
+            inner,
+            stream_id,
+            available_data,
+        } = self;
+        let stream_id = *stream_id;
+        let (beg, end) = inner.map_read(stream_id)?;
+        let nbytes = unsafe { byte_offset_from(beg, end) };
+        if nbytes > 0 {
+            log::trace!(
+                "[stream {}] ACQUIRED {:p}-{:p}:{} bytes",
+                stream_id,
+                beg,
+                end,
+                nbytes
+            )
+        };
+        if nbytes > 0 {
+            *available_data = Some(Python::with_gil(|py| {
+                Py::new(
+                    py,
+                    AvailableData {
+                        inner: Some(Arc::new(RawAvailableData {
+                            runtime: self.inner.clone(),
+                            beg: NonNull::new(beg).ok_or(anyhow!("Expected non-null buffer"))?,
+                            end: NonNull::new(end).ok_or(anyhow!("Expected non-null buffer"))?,
+                            stream_id,
+                            consumed_bytes: None,
+                        })),
+                    },
+                )
+            })?);
+        }
+        return Ok(self.available_data.clone());
     }
-    fn __next__(&mut self) -> Option<VideoFrame> {
-        self.next()
+
+    fn __exit__(&mut self, _exc_type: &PyAny, _exc_value: &PyAny, _traceback: &PyAny) {
+        Python::with_gil(|py| {
+            if let Some(a) = &self.available_data {
+                a.as_ref(py).borrow_mut().invalidate()
+            };
+        });
     }
 }
 
-impl Iterator for VideoFrameIterator {
+struct VideoFrameIteratorInner {
+    store: Arc<RawAvailableData>,
+    cur: Mutex<NonNull<capi::VideoFrame>>,
+    end: NonNull<capi::VideoFrame>,
+}
+
+unsafe impl Send for VideoFrameIteratorInner {}
+
+impl Iterator for VideoFrameIteratorInner {
     type Item = VideoFrame;
 
     fn next(&mut self) -> Option<Self::Item> {
@@ -349,6 +396,29 @@ impl Iterator for VideoFrameIterator {
     }
 }
 
+#[pyclass]
+struct VideoFrameIterator {
+    inner: Option<VideoFrameIteratorInner>,
+}
+
+#[pymethods]
+impl VideoFrameIterator {
+    fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
+        slf
+    }
+    fn __next__(&mut self) -> Option<VideoFrame> {
+        self.next()
+    }
+}
+
+impl Iterator for VideoFrameIterator {
+    type Item = VideoFrame;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        self.inner.as_mut().and_then(|it| it.next())
+    }
+}
+
 #[pyclass]
 #[derive(Debug, Default, Clone, Copy, Serialize, Deserialize)]
 pub(crate) struct VideoFrameTimestamps {
@@ -489,6 +559,5 @@ impl VideoFrame {
     }
 }
 
-// TODO: Can't ensure the output array doesn't outlive the available data
 // TODO: Is everything really Send
 // TODO: mark iterable and videoframe as things that can't be shared across threads