Consolidate information into ScanData, remove ad-hoc differences for specific error in _run_scan

src/openflexure_microscope_server/scan_directories.py

@@ -1,6 +1,6 @@
 """Functionality to manage file system operations for scan directories."""
 
-from typing import Optional, Any
+from typing import Optional, Any, Mapping
 import os
 import re
 import shutil
@@ -54,6 +54,7 @@ class ScanData(BaseModel):
     # TODO: Docs for each variable
     # TODO: Make note about timestamp format. Is it too ambiguous?
     scan_name: str
+    starting_position: Mapping[str, int]
     overlap: float
     max_dist: int
     dx: int
@@ -67,18 +68,16 @@ class ScanData(BaseModel):
     # just for correlation stitching
     stitch_resize: float
     save_resolution: tuple[int, int]
-    final_image_count: Optional[int] = None
+    image_count: int = 0
     duration: Optional[timedelta] = None
     scan_result: Optional[str] = None
 
-    def set_final_data(self, result: str, final_image_count: int):
+    def set_final_data(self, result: str):
         """Set the final data for the scan, scan duration is automatically calculated.
 
         :param result: A string describing the result.
-        :param final_image_count: The total number of images captures.
         """
         self.duration = datetime.now() - self.start_time
-        self.final_image_count = final_image_count
         self.scan_result = result
 
     @field_validator("start_time", mode="before")
@@ -118,7 +117,7 @@ class ScanData(BaseModel):
         secs = int((total_secs % 60))
         return f"{hrs}:{mins:02}:{secs:02}"
 
-    @field_validator("final_image_count", "scan_result", mode="before")
+    @field_validator("scan_result", mode="before")
     @classmethod
     def parse_unknown_as_none(cls, value: Optional[str | int]) -> Optional[str | int]:
         """Validate the string "Unknown" as None."""
@@ -126,7 +125,7 @@ class ScanData(BaseModel):
             return None
         return value
 
-    @field_serializer("final_image_count", "scan_result")
+    @field_serializer("scan_result")
     def serialize_none_as_unknown(self, value: Optional[str | int]) -> str | int:
         """Serialise None as "Unknown" for a more human readable result."""
         return "Unknown" if value is None else value

src/openflexure_microscope_server/things/smart_scan.py

@@ -7,7 +7,7 @@ It also controls external processes for live stitching composite images, and
 the creation of the final stitched images.
 """
 
-from typing import Optional, Mapping
+from typing import Optional
 import threading
 import os
 import time
@@ -88,20 +88,16 @@ class SmartScanThing(lt.Thing):
         # when the `sample_scan` lt.thing_action is called. It is saved as
         # private class variable along with many others here.
         # Access to these variables requires a scan to be running,
-        # any method that calls these should be decorrected with
+        # any method that calls these should be decorated with
         # @_scan_running
         self._scan_logger: Optional[lt.deps.InvocationLogger] = None
         self._cancel: Optional[lt.deps.CancelHook] = None
         self._autofocus: Optional[AutofocusDep] = None
         self._stage: Optional[StageDep] = None
         self._cam: Optional[CamDep] = None
-        self._metadata_getter: Optional[lt.deps.GetThingStates] = None
         self._csm: Optional[CSMDep] = None
+
         self._ongoing_scan: Optional[scan_directories.ScanDirectory] = None
-        # TODO see if starting position can go into ScanData
-        self._starting_position: Optional[Mapping[str, int]] = None
-        self._scan_images_taken: Optional[int] = None
-        # TODO Scan data is a dict during refactoring, should become a dataclass
         self._scan_data: Optional[scan_directories.ScanData] = None
         self._preview_stitcher: Optional[stitching.PreviewStitcher] = None
 
@@ -113,7 +109,6 @@ class SmartScanThing(lt.Thing):
         autofocus: AutofocusDep,
         stage: StageDep,
         cam: CamDep,
-        metadata_getter: lt.deps.GetThingStates,
         csm: CSMDep,
         scan_name: str = "",
     ):
@@ -133,22 +128,14 @@ class SmartScanThing(lt.Thing):
         self._autofocus = autofocus
         self._stage = stage
         self._cam = cam
-        # TODO check if metadata_getter this can be removed without error?
-        self._metadata_getter = metadata_getter
         self._csm = csm
-        self._scan_images_taken = 0
+        # Confirm scan data is None as start of scan.
-
-        # Set _scan_data to None. This is needed just in case an exception is raised
-        # before _run_scan (which sets the real data). As we check this in the `except`
         self._scan_data = None
-
         try:
             self._check_background_and_csm_set()
             self._ongoing_scan = self._scan_dir_manager.new_scan_dir(scan_name)
             self._latest_scan_name = self._ongoing_scan.name
             self._autofocus.looping_autofocus(dz=self.autofocus_dz, start="centre")
-            # record starting position so we can return there
-            self._starting_position = self._stage.position
             self._run_scan()
         except Exception as e:
             # If _scan_data is set then scan started
@@ -156,6 +143,9 @@ class SmartScanThing(lt.Thing):
                 self._return_to_starting_position()
                 if not isinstance(e, scan_directories.NotEnoughFreeSpaceError):
                     # Don't stitch if drive is full (already logged)
+                    self._scan_logger.info(
+                        "Attempting to stitch and archive the images acquired so far."
+                    )
                     self._perform_final_stitch()
             # Error must be raised so UI gives correct output
             raise e
@@ -166,10 +156,8 @@ class SmartScanThing(lt.Thing):
             self._autofocus = None
             self._stage = None
             self._cam = None
-            self._metadata_getter = None
             self._csm = None
             self._ongoing_scan = None
-            self._scan_images_taken = None
             self._scan_data = None
             self._scan_lock.release()
             # Ensure any PreviewStitcher created cannot be reused.
@@ -275,6 +263,8 @@ class SmartScanThing(lt.Thing):
     @_scan_running
     def _collect_scan_data(self) -> scan_directories.ScanData:
         """Collect and return the data for this scan so it cannot be changed mid-scan."""
+        # Record starting position so it can be returned to at end of scan.
+        starting_position = self._stage.position
         overlap = self.overlap
         dx, dy = self._calc_displacement_from_test_image(overlap)
         stitch_resize = stitching.STITCHING_RESOLUTION[0] / self.save_resolution[0]
@@ -300,6 +290,7 @@ class SmartScanThing(lt.Thing):
         # Fix scan parameters in case UI is updated during scan.
         return scan_directories.ScanData(
             scan_name=self._ongoing_scan.name,
+            starting_position=starting_position,
             overlap=overlap,
             max_dist=self.max_range,
             dx=dx,
@@ -320,10 +311,7 @@ class SmartScanThing(lt.Thing):
         Takes scan_result, a string that is either "success", "cancelled by user",
         or the error that ended the scan.
         """
-        self._scan_data.set_final_data(
+        self._scan_data.set_final_data(result=scan_result)
-            result=scan_result,
-            final_image_count=self._scan_images_taken,
-        )
         self._ongoing_scan.save_scan_data(self._scan_data)
 
     @_scan_running
@@ -331,7 +319,7 @@ class SmartScanThing(lt.Thing):
         """Manage the stitching threads, starting them if needed and not already running."""
         # Assume 4 images means at least one offset in x and y, making the stitching
         # well constrained.
-        if self._scan_images_taken > 3:
+        if self._scan_data.image_count > 3:
             if not self._preview_stitcher.running:
                 self._preview_stitcher.start()
 
@@ -353,10 +341,6 @@ class SmartScanThing(lt.Thing):
                 correlation_resize=self._scan_data.stitch_resize,
             )
 
-            if self._scan_images_taken != 0:
-                msg = "_scan_images_taken should be zero before starting scanning"
-                raise RuntimeError(msg)
-
             # This is the main loop of the scan!
             self._main_scan_loop()
             self._save_final_scan_data(scan_result="success")
@@ -366,24 +350,19 @@ class SmartScanThing(lt.Thing):
             self._cancel.clear()
             self._scan_logger.info("Stopping scan because it was cancelled.")
             self._save_final_scan_data(scan_result="cancelled by user")
-        except scan_directories.NotEnoughFreeSpaceError as e:
-            self._save_final_scan_data(scan_result=f"NotEnoughFreeSpaceError: {e}")
-            self._scan_logger.error(
-                f"Stopping scan to avoid filling up the disk: {e}",
-                exc_info=e,
-            )
-            raise e
         except Exception as e:
             err_name = type(e).__name__
-            self._save_final_scan_data(scan_result=f"{err_name}: {e}")
+            if self._scan_data is not None:
+                self._save_final_scan_data(scan_result=f"{err_name}: {e}")
             self._scan_logger.error(
-                f"The scan stopped because of an error: {e} "
+                f"The scan stopped because of an error: {e}",
-                "Attempting to stitch and archive the images acquired so far.",
                 exc_info=e,
             )
             raise e
         finally:
-            # Don't set Preview Stitcher to None yet. It is used by _perform_final_scan.
+            # Don't set Preview Stitcher to None yet. It is used by
+            # _perform_final_stitch, which may also be run after this function completes
+            # if it ended due to an exception.
 
             # Start streaming in the default resolution again as soon as possible
             self._cam.start_streaming()
@@ -459,7 +438,7 @@ class SmartScanThing(lt.Thing):
             )
 
             # increment capture counter as thread has completed
-            self._scan_images_taken += 1
+            self._scan_data.image_count += 1
             # Add it to the incremental zip
             self._ongoing_scan.zip_files()
 
@@ -467,15 +446,15 @@ class SmartScanThing(lt.Thing):
     def _return_to_starting_position(self):
         """Return to the initial scan position, if set."""
         self._scan_logger.info("Returning to starting position.")
-        if self._starting_position is not None:
+        if self._scan_data is not None:
             self._stage.move_absolute(
-                **self._starting_position, block_cancellation=True
+                **self._scan_data.starting_position, block_cancellation=True
             )
 
     @_scan_running
     def _perform_final_stitch(self):
         """Update the scan zip and perform final stitch of the data."""
-        if self._scan_images_taken <= 3:
+        if self._scan_data.image_count <= 3:
             self._scan_logger.info("Not performing a stitch as 3 or fewer images taken")
             return
 

tests/test_scan_data.py

@@ -35,6 +35,7 @@ def _fake_scan_data(**kwargs) -> ScanData:
     """
     data_dict = {
         "scan_name": "fake_scan_0001",
+        "starting_position": {"x": 123, "y": 456, "z": 789},
         "overlap": 0.1,
         "max_dist": 100000,
         "dx": 100,
@@ -56,17 +57,19 @@ def test_set_final_data():
     """Check that adding final data to a ScanData object works as expected."""
     scan_data = _fake_scan_data()
 
-    assert scan_data.final_image_count is None
+    assert scan_data.image_count == 0
     assert scan_data.duration is None
     assert scan_data.scan_result is None
 
+    # Quickly take 123 images!
+    scan_data.image_count += 123
     # Should set duration based of finishing at datetime.now()
-    scan_data.set_final_data(result="Success", final_image_count=123)
+    scan_data.set_final_data(result="Success")
     expected_duration = datetime.now() - MOCK_START_TIME
     expected_duration_s = expected_duration.total_seconds()
     scan_duration_s = scan_data.duration.total_seconds()
 
-    assert scan_data.final_image_count == 123
+    assert scan_data.image_count == 123
     assert expected_duration_s - 1 < scan_duration_s < expected_duration_s + 1
     assert scan_data.scan_result == "Success"
 
@@ -77,15 +80,16 @@ def test_custom_serialisation():
     # Serialise to string then load directly as json
     scan_data_dict = json.loads(scan_data.model_dump_json())
     assert scan_data_dict["start_time"] == "11_00_00-25_12_2024"
-    assert scan_data_dict["final_image_count"] == "Unknown"
+    assert scan_data_dict["image_count"] == 0
     assert scan_data_dict["duration"] == "Unknown"
     assert scan_data_dict["scan_result"] == "Unknown"
 
-    scan_data.set_final_data(result="Success", final_image_count=123)
+    scan_data.image_count += 123
+    scan_data.set_final_data(result="Success")
     # Can't mock datetime.now as datetime is immutable. So just replace duraction
     scan_data.duration = MOCK_END_TIME - scan_data.start_time
     scan_data_dict = json.loads(scan_data.model_dump_json())
-    assert scan_data_dict["final_image_count"] == 123
+    assert scan_data_dict["image_count"] == 123
     assert scan_data_dict["duration"] == "1:29:03"
     assert scan_data_dict["scan_result"] == "Success"
 
@@ -106,7 +110,8 @@ def test_round_trip_finalised():
     """Check that finalised ScanData can be serialised and deserialised."""
     scan_data = _fake_scan_data()
     # Finalise the data.
-    scan_data.set_final_data(result="Success", final_image_count=123)
+    scan_data.image_count += 123
+    scan_data.set_final_data(result="Success")
 
     scan_data_dict = json.loads(scan_data.model_dump_json())
     scan_data_reloaded = ScanData(**scan_data_dict)

tests/test_smart_scan.py

@@ -177,7 +177,6 @@ def _run_only_outer_scan(adjust_initial_state: Optional[Callable] = None):
     af_mock = MockAutoFocusThing()
     stage_mock = MockStageThing()
     cam_mock = MockCameraThing()
-    meta_mock = 5  # not called
     csm_mock = MockCSMThing()
 
     class MockedSmartScanThing(SmartScanThing):
@@ -196,9 +195,7 @@ def _run_only_outer_scan(adjust_initial_state: Optional[Callable] = None):
             assert self._autofocus is af_mock
             assert self._stage is stage_mock
             assert self._cam is cam_mock
-            assert self._metadata_getter is meta_mock
             assert self._csm is csm_mock
-            assert self._scan_images_taken == 0
 
     # mock smart scan thing
     mock_ss_thing = MockedSmartScanThing(SCAN_DIR)
@@ -214,7 +211,6 @@ def _run_only_outer_scan(adjust_initial_state: Optional[Callable] = None):
             autofocus=af_mock,
             stage=stage_mock,
             cam=cam_mock,
-            metadata_getter=meta_mock,
             csm=csm_mock,
             scan_name="FooBar",
         )
@@ -228,9 +224,7 @@ def _run_only_outer_scan(adjust_initial_state: Optional[Callable] = None):
     assert mock_ss_thing._autofocus is None
     assert mock_ss_thing._stage is None
     assert mock_ss_thing._cam is None
-    assert mock_ss_thing._metadata_getter is None
     assert mock_ss_thing._csm is None
-    assert mock_ss_thing._scan_images_taken is None
 
     # Return the mock thing for further state testing, and the
     # exec_info of any uncaught exceptions that were raised
@@ -260,10 +254,16 @@ def test_outer_scan_wo_sample_skip():
     assert mock_ss_thing.mock_call_count["_run_scan"] == 1
 
 
+MOCK_SCAN_NAME = "test_name_0001"
+MOCK_SCAN_DIR = "scans/test_name_0001/images/"
+MOCK_START_POS = {"x": 123, "y": 456, "z": 789}
+
+
 def _expected_scan_data():
     """Return the expected ScanData object for a SmartScan with default properties."""
     expected_dict = {
-        "scan_name": "test_name_0001",
+        "scan_name": MOCK_SCAN_NAME,
+        "starting_position": MOCK_START_POS,
         "overlap": 0.45,
         "max_dist": 45000,
         "dx": 100,
@@ -287,11 +287,13 @@ def scan_thing_mocked_for_scan_data(smart_scan_thing, mocker):
         smart_scan_thing, "_calc_displacement_from_test_image", return_value=[100, 100]
     )
     mock_ongoing_scan = mocker.Mock()
-    type(mock_ongoing_scan).name = mocker.PropertyMock(return_value="test_name_0001")
+    type(mock_ongoing_scan).name = mocker.PropertyMock(return_value=MOCK_SCAN_NAME)
-    type(mock_ongoing_scan).images_dir = mocker.PropertyMock(
+    type(mock_ongoing_scan).images_dir = mocker.PropertyMock(return_value=MOCK_SCAN_DIR)
-        return_value="scans/test_name_0001/images/"
+    mock_stage = mocker.Mock()
-    )
+    type(mock_stage).position = mocker.PropertyMock(return_value=MOCK_START_POS)
+
     smart_scan_thing._ongoing_scan = mock_ongoing_scan
+    smart_scan_thing._stage = mock_stage
     return smart_scan_thing
 
 
@@ -313,7 +315,7 @@ def test_save_final_scan_data(scan_thing_mocked_for_scan_data):
     scan_thing = scan_thing_mocked_for_scan_data
 
     scan_thing._scan_data = scan_thing._collect_scan_data()
-    scan_thing._scan_images_taken = 44
+    scan_thing._scan_data.image_count = 44
     scan_thing._save_final_scan_data("Mocked!")
     # _ongoing_scan is a mock so we can check that save_scan data was called and get
     # the value
@@ -321,7 +323,7 @@ def test_save_final_scan_data(scan_thing_mocked_for_scan_data):
     final_data = scan_thing._ongoing_scan.save_scan_data.call_args[0][0]
     assert isinstance(final_data, ScanData)
     assert final_data.scan_result == "Mocked!"
-    assert final_data.final_image_count == 44
+    assert final_data.image_count == 44
     assert final_data.duration.total_seconds() < 1
 
 
@@ -335,7 +337,6 @@ def scan_thing_mocked_for_run_scan(scan_thing_mocked_for_scan_data, mocker):
     """
     scan_thing = scan_thing_mocked_for_scan_data
     scan_thing._cam = MockCameraThing()
-    scan_thing._scan_images_taken = 0
     mocker.patch.object(scan_thing, "_cancel")
     mocker.patch.object(scan_thing, "_main_scan_loop")
     mocker.patch.object(scan_thing, "_return_to_starting_position")
@@ -364,7 +365,7 @@ def check_run_scan(scan_thing, caplog, expected_exception=None):
         with pytest.raises(expected_exception), caplog.at_level(logging.WARNING):
             scan_thing._scan_data = scan_thing._run_scan()
     # The preview stitcher object should still exist. And images dir should be set.
-    assert scan_thing._preview_stitcher.images_dir == "scans/test_name_0001/images/"
+    assert scan_thing._preview_stitcher.images_dir == MOCK_SCAN_DIR
 
     final_scan_data = scan_thing._ongoing_scan.save_scan_data.call_args[0][0]
     calls = {
@@ -396,34 +397,6 @@ def test_run_scan(scan_thing_mocked_for_run_scan, caplog):
     assert calls == expected_calls_numbers
 
 
-def test_run_scan_wrong_image_value(scan_thing_mocked_for_run_scan, caplog):
-    """Check correct methods called if _scan_images_taken doesn't start at 0.
-
-    This should never happen but shows a good clear path for an error before
-    _main_scan_loop starts.
-    """
-    scan_thing = scan_thing_mocked_for_run_scan
-    scan_thing._scan_images_taken = 2
-
-    result, logs, calls = check_run_scan(scan_thing, caplog, RuntimeError)
-
-    assert result.startswith("RuntimeError:")
-    assert len(logs) == 1
-    assert logs[0].levelno == logging.ERROR
-
-    # Main loop not run, nor are return to start, final stitch, or purging of empty
-    # scans. Save scan data is still called twice
-    expected_calls_numbers = {
-        "cam_start_streaming_calls": 2,
-        "main_scan_loop_calls": 0,
-        "return_to_start_calls": 0,
-        "perform_final_stitch_calls": 0,
-        "purge_empty_scans_calls": 0,
-        "save_scan_data_calls": 2,
-    }
-    assert calls == expected_calls_numbers
-
-
 def test_run_scan_err_in_main_loop(scan_thing_mocked_for_run_scan, caplog, mocker):
     """Check correct methods called if main_loop errors."""
     scan_thing = scan_thing_mocked_for_run_scan