Support empty AvailableData structs

Cargo.toml

@@ -1,7 +1,7 @@
 [package]
 name = "acquire-imaging"
 authors = ["Nathan Clack <[email protected]>"]
-version = "0.3.0"
+version = "0.3.0-rc1"
 edition = "2021"
 
 [lib]

drivers.json

@@ -1,6 +1,6 @@
 {
   "acquire-driver-common": "0.1.6",
-  "acquire-driver-zarr": "0.1.7",
+  "acquire-driver-zarr": "0.1.8",
   "acquire-driver-egrabber": "0.1.5",
   "acquire-driver-hdcam": "0.1.7",
   "acquire-driver-spinnaker": "0.1.1"

python/acquire/acquire.pyi

@@ -14,7 +14,7 @@ from numpy.typing import NDArray
 
 @final
 class AvailableDataContext:
-    def __enter__(self) -> Optional[AvailableData]: ...
+    def __enter__(self) -> AvailableData: ...
     def __exit__(
         self, exc_type: Any, exc_value: Any, traceback: Any
     ) -> None: ...

src/runtime.rs

@@ -196,7 +196,7 @@ impl Runtime {
         Ok(AvailableDataContext {
             inner: self.inner.clone(),
             stream_id,
-            available_data: None,
+            available_data: Python::with_gil(|py| Py::new(py, AvailableData { inner: None }))?,
         })
     }
 }
@@ -273,7 +273,7 @@ impl Drop for RawAvailableData {
 
 #[pyclass]
 pub(crate) struct AvailableData {
-    inner: Option<Arc<RawAvailableData>>,
+    inner: Option<RawAvailableData>,
 }
 
 #[pymethods]
@@ -290,7 +290,6 @@ impl AvailableData {
         VideoFrameIterator {
             inner: if let Some(frames) = &self.inner {
                 Some(VideoFrameIteratorInner {
-                    store: frames.clone(),
                     cur: Mutex::new(frames.beg),
                     end: frames.end,
                 })
@@ -315,12 +314,12 @@ impl AvailableData {
 pub(crate) struct AvailableDataContext {
     inner: Arc<RawRuntime>,
     stream_id: u32,
-    available_data: Option<Py<AvailableData>>,
+    available_data: Py<AvailableData>,
 }
 
 #[pymethods]
 impl AvailableDataContext {
-    fn __enter__(&mut self) -> PyResult<Option<Py<AvailableData>>> {
+    fn __enter__(&mut self) -> PyResult<Py<AvailableData>> {
         let AvailableDataContext {
             inner,
             stream_id,
@@ -329,45 +328,39 @@ impl AvailableDataContext {
         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
+
+        log::trace!(
+            "[stream {}] ACQUIRED {:p}-{:p}:{} bytes",
+            stream_id,
+            beg,
+            end,
+            nbytes
+        );
+        *available_data = Python::with_gil(|py| {
+            Py::new(
+                py,
+                AvailableData {
+                    inner: Some(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,
+                    }),
+                },
             )
-        };
-        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 __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()
-            };
+            (&self.available_data).as_ref(py).borrow_mut().invalidate();
         });
     }
 }
 
 struct VideoFrameIteratorInner {
-    store: Arc<RawAvailableData>,
     cur: Mutex<NonNull<capi::VideoFrame>>,
     end: NonNull<capi::VideoFrame>,
 }
@@ -380,10 +373,7 @@ impl Iterator for VideoFrameIteratorInner {
     fn next(&mut self) -> Option<Self::Item> {
         let mut cur = self.cur.lock();
         if *cur < self.end {
-            let out = VideoFrame {
-                _store: self.store.clone(),
-                cur: *cur,
-            };
+            let out = VideoFrame { cur: *cur };
 
             let c = cur.as_ptr();
             let o = unsafe { (c as *const u8).offset((*c).bytes_of_frame as _) }
@@ -503,7 +493,6 @@ impl IntoDimension for capi::ImageShape {
 
 #[pyclass]
 pub(crate) struct VideoFrame {
-    _store: Arc<RawAvailableData>,
     cur: NonNull<capi::VideoFrame>,
 }
 

tests/test_basic.py

@@ -96,9 +96,8 @@ def test_repeat_acq(runtime: Runtime):
     runtime.start()
     while True:
         with runtime.get_available_data(0) as a:
-            if a:
-                logging.info(f"Got {a.get_frame_count()}")
-                break
+            logging.info(f"Got {a.get_frame_count()}")
+            break
         if a:
             assert a.get_frame_count() == 0
             assert next(a.frames()) is None
@@ -107,12 +106,11 @@ def test_repeat_acq(runtime: Runtime):
     runtime.start()
     while True:
         with runtime.get_available_data(0) as a:
-            if a:
-                logging.info(f"Got {a.get_frame_count()}")
-                break
-        if a:
-            assert a.get_frame_count() == 0
-            assert next(a.frames()) is None
+            logging.info(f"Got {a.get_frame_count()}")
+            break
+
+        assert a.get_frame_count() == 0
+        assert next(a.frames()) is None
     runtime.stop()
     # TODO: (nclack) assert 1 more acquired frame. stop cancels and waits.
 
@@ -131,7 +129,7 @@ def test_repeat_with_no_stop(runtime: Runtime):
     # wait for 1 frame
     while True:
         with runtime.get_available_data(0) as a:
-            if a:
+            if a.get_frame_count() > 0:
                 logging.info(f"Got {a.get_frame_count()} frame")
                 break
     # acq is still on going here
@@ -181,17 +179,16 @@ def took_too_long():
     while nframes < p.video[0].max_frame_count and not took_too_long():
         clock = time.time()
         with runtime.get_available_data(0) as a:
-            if a:
-                packet = a.get_frame_count()
-                for f in a.frames():
-                    logging.info(
-                        f"{f.data().shape} {f.data()[0][0][0][0]} "
-                        + f"{f.metadata()}"
-                    )
-                nframes += packet
+            packet = a.get_frame_count()
+            for f in a.frames():
                 logging.info(
-                    f"frame count: {nframes} - frames in packet: {packet}"
+                    f"{f.data().shape} {f.data()[0][0][0][0]} "
+                    + f"{f.metadata()}"
                 )
+            nframes += packet
+            logging.info(
+                f"frame count: {nframes} - frames in packet: {packet}"
+            )
 
         elapsed = time.time() - clock
         sleep(max(0, 0.1 - elapsed))
@@ -231,8 +228,7 @@ def test_change_filename(runtime: Runtime):
         runtime.start()
         while nframes < p.video[0].max_frame_count:
             with runtime.get_available_data(0) as packet:
-                if packet:
-                    nframes += packet.get_frame_count()
+                nframes += packet.get_frame_count()
         logging.info("Stopping")
         runtime.stop()
 
@@ -257,8 +253,7 @@ def test_write_external_metadata_to_tiff(
     runtime.start()
     while nframes < p.video[0].max_frame_count:
         with runtime.get_available_data(0) as packet:
-            if packet:
-                nframes += packet.get_frame_count()
+            nframes += packet.get_frame_count()
     runtime.stop()
 
     # Check that the written tif has the expected structure
@@ -321,20 +316,17 @@ def is_not_done() -> bool:
     while is_not_done():
         if nframes[stream_id] < p.video[stream_id].max_frame_count:
             with runtime.get_available_data(stream_id) as packet:
-                if packet:
-                    n = packet.get_frame_count()
-                    for i, frame in enumerate(packet.frames()):
-                        expected_frame_id = nframes[stream_id] + i
-                        assert (
-                            frame.metadata().frame_id == expected_frame_id
-                        ), (
-                            "frame id's didn't match "
-                            + f"({frame.metadata().frame_id}"
-                            + f"!={expected_frame_id})"
-                            + f" [stream {stream_id} nframes {nframes}]"
-                        )
-                    nframes[stream_id] += n
-                    logging.debug(f"NFRAMES {nframes}")
+                n = packet.get_frame_count()
+                for i, frame in enumerate(packet.frames()):
+                    expected_frame_id = nframes[stream_id] + i
+                    assert frame.metadata().frame_id == expected_frame_id, (
+                        "frame id's didn't match "
+                        + f"({frame.metadata().frame_id}"
+                        + f"!={expected_frame_id})"
+                        + f" [stream {stream_id} nframes {nframes}]"
+                    )
+                nframes[stream_id] += n
+                logging.debug(f"NFRAMES {nframes}")
 
         stream_id = (stream_id + 1) % 2
     logging.info("Stopping")
@@ -362,9 +354,9 @@ def test_abort(runtime: Runtime):
 
     while True:
         with runtime.get_available_data(0) as packet:
-            if packet:
-                nframes += packet.get_frame_count()
-            else:
+            frame_count = packet.get_frame_count()
+            nframes += frame_count
+            if frame_count == 0:
                 break
 
     logging.debug(
@@ -383,7 +375,7 @@ def wait_for_data(
     elapsed = timedelta()
     while elapsed < timeout:
         with runtime.get_available_data(stream_id) as packet:
-            if packet:
+            if packet.get_frame_count() > 0:
                 frames = list(packet.frames())
                 return (len(frames), frames[0].metadata().frame_id)
         sleep(sleep_duration.total_seconds())
@@ -419,7 +411,7 @@ def test_execute_trigger(runtime: Runtime):
 
     # No triggers yet, so no data.
     with runtime.get_available_data(0) as packet:
-        assert packet is None
+        assert packet.get_frame_count() == 0
 
     # Snap a few individual frames
     for i in range(p.video[0].max_frame_count):

tests/test_zarr.py

@@ -47,8 +47,7 @@ def test_write_external_metadata_to_zarr(
     runtime.start()
     while nframes < p.video[0].max_frame_count:
         with runtime.get_available_data(0) as packet:
-            if packet:
-                nframes += packet.get_frame_count()
+            nframes += packet.get_frame_count()
     runtime.stop()
 
     assert p.video[0].storage.settings.filename