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

This commit is contained in:
Julian Stirling 2025-08-04 10:09:39 +01:00
parent 82691a6b78
commit 4e9b07c78f
4 changed files with 56 additions and 100 deletions

View file

@ -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

View file

@ -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
# 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`
# Confirm scan data is None as start of scan.
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(
result=scan_result,
final_image_count=self._scan_images_taken,
)
self._scan_data.set_final_data(result=scan_result)
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} "
"Attempting to stitch and archive the images acquired so far.",
f"The scan stopped because of an error: {e}",
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

View file

@ -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)

View file

@ -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).images_dir = mocker.PropertyMock(
return_value="scans/test_name_0001/images/"
)
type(mock_ongoing_scan).name = mocker.PropertyMock(return_value=MOCK_SCAN_NAME)
type(mock_ongoing_scan).images_dir = mocker.PropertyMock(return_value=MOCK_SCAN_DIR)
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