diff --git a/bec_widgets/cli/client.py b/bec_widgets/cli/client.py index 539158c9..92c108e8 100644 --- a/bec_widgets/cli/client.py +++ b/bec_widgets/cli/client.py @@ -2360,6 +2360,7 @@ class Heatmap(RPCBase): oversampling_factor: "float | None" = None, lock_aspect_ratio: "bool | None" = None, show_config_label: "bool | None" = None, + scan_id: "str | None" = None, reload: "bool" = False, ): """ @@ -2379,9 +2380,16 @@ class Heatmap(RPCBase): oversampling_factor (float | None): Factor to oversample the grid resolution. lock_aspect_ratio (bool | None): Whether to lock the aspect ratio of the image. show_config_label (bool | None): Whether to show the configuration label in the heatmap. + scan_id (str | None): Optional scan ID to fetch from history instead of using the live/latest scan. reload (bool): Whether to reload the heatmap with new data. """ + @rpc_call + def update_with_scan_history(self, scan_id: "str | None" = None): + """ + Update the heatmap with a scan fetched from history. + """ + @property @rpc_call def device_x(self) -> "str": diff --git a/bec_widgets/widgets/plots/heatmap/heatmap.py b/bec_widgets/widgets/plots/heatmap/heatmap.py index e3338ee5..c1b2a67d 100644 --- a/bec_widgets/widgets/plots/heatmap/heatmap.py +++ b/bec_widgets/widgets/plots/heatmap/heatmap.py @@ -214,6 +214,7 @@ class Heatmap(ImageBase): "remove_roi", "rois", "plot", + "update_with_scan_history", # Device properties "device_x", "device_x.setter", @@ -338,6 +339,7 @@ class Heatmap(ImageBase): oversampling_factor: float | None = None, lock_aspect_ratio: bool | None = None, show_config_label: bool | None = None, + scan_id: str | None = None, reload: bool = False, ): """ @@ -357,6 +359,7 @@ class Heatmap(ImageBase): oversampling_factor (float | None): Factor to oversample the grid resolution. lock_aspect_ratio (bool | None): Whether to lock the aspect ratio of the image. show_config_label (bool | None): Whether to show the configuration label in the heatmap. + scan_id (str | None): Optional scan ID to fetch from history instead of using the live/latest scan. reload (bool): Whether to reload the heatmap with new data. """ if validate_bec: @@ -418,6 +421,10 @@ class Heatmap(ImageBase): self.main_image.clear() self.update_labels() + if scan_id is not None: + self.update_with_scan_history(scan_id=scan_id) + return + self._fetch_running_scan() self.sync_signal_update.emit() @@ -431,6 +438,50 @@ class Heatmap(ImageBase): self.scan_id = self.client.history._scan_ids[-1] self.old_scan_id = None + def get_history_scan_item(self, scan_id: str | None = None): + """Fetch a scan item from history, defaulting to the latest historical scan.""" + if self.client.history is None or len(self.client.history) == 0: + logger.info("No scans executed so far. Cannot fetch scan history.") + return None + + if scan_id is None: + return self.client.history[-1] + return self.client.history.get_by_scan_id(scan_id) + + @SafeSlot(str) + @SafeSlot() + def update_with_scan_history(self, scan_id: str | None = None): + """Update the heatmap with a scan fetched from history.""" + scan_item = self.get_history_scan_item(scan_id=scan_id) + if scan_item is None: + return + + if scan_id is not None: + target_scan_id = scan_id + elif hasattr(scan_item, "metadata"): + target_scan_id = scan_item.metadata["bec"]["scan_id"] + else: + target_scan_id = scan_item.scan_id + + if target_scan_id != self.scan_id: + self._invalidate_interpolation_generation() + self._grid_index = None + self.main_image.clear() + self.status_message = None + + self.scan_item = scan_item + if self.scan_item is None: + return + + if scan_id is not None: + self.scan_id = scan_id + elif hasattr(self.scan_item, "metadata"): + self.scan_id = self.scan_item.metadata["bec"]["scan_id"] + else: + self.scan_id = self.scan_item.scan_id + + self.sync_signal_update.emit() + def update_labels(self): """ Update the labels of the x, y, and z axes. diff --git a/tests/unit_tests/test_heatmap_widget.py b/tests/unit_tests/test_heatmap_widget.py index 38504ba4..730037e4 100644 --- a/tests/unit_tests/test_heatmap_widget.py +++ b/tests/unit_tests/test_heatmap_widget.py @@ -37,6 +37,65 @@ def test_heatmap_plot(heatmap_widget): assert heatmap_widget._image_config.device_z.device == "bpm4i" +def test_heatmap_plot_with_scan_id_uses_history(heatmap_widget): + history_scan = mock.MagicMock() + history_scan.scan_id = "scan-123" + + with ( + mock.patch.object( + heatmap_widget, + "update_with_scan_history", + wraps=heatmap_widget.update_with_scan_history, + ) as update_history_mock, + mock.patch.object(heatmap_widget, "get_history_scan_item", return_value=history_scan), + ): + heatmap_widget.plot( + device_x="samx", + device_y="samy", + device_z="bpm4i", + signal_x="samx", + signal_y="samy", + signal_z="bpm4i", + scan_id="scan-123", + ) + + update_history_mock.assert_called_once_with(scan_id="scan-123") + assert heatmap_widget.scan_item is history_scan + assert heatmap_widget.scan_id == "scan-123" + + +def test_heatmap_update_with_scan_history_resets_cached_image_state(heatmap_widget): + history_scan = mock.MagicMock() + history_scan.scan_id = "scan-456" + + heatmap_widget.scan_id = "scan-123" + heatmap_widget._grid_index = 17 + heatmap_widget.status_message = mock.MagicMock() + + with ( + mock.patch.object(heatmap_widget, "get_history_scan_item", return_value=history_scan), + mock.patch.object(heatmap_widget.main_image, "clear") as clear_mock, + ): + heatmap_widget.update_with_scan_history(scan_id="scan-456") + + clear_mock.assert_called_once() + assert heatmap_widget._grid_index is None + assert heatmap_widget.status_message is None + assert heatmap_widget.scan_item is history_scan + assert heatmap_widget.scan_id == "scan-456" + + +def test_heatmap_on_scan_status_resets_after_history_scan_selection(heatmap_widget): + heatmap_widget.scan_id = "scan-123" + scan_msg = messages.ScanStatusMessage(scan_id="live-scan", status="open", metadata={}, info={}) + + with mock.patch.object(heatmap_widget, "reset") as reset_mock: + heatmap_widget.on_scan_status(scan_msg.content, scan_msg.metadata) + + reset_mock.assert_called_once() + assert heatmap_widget.scan_id == "live-scan" + + def test_heatmap_on_scan_status_no_scan_id(heatmap_widget): scan_msg = messages.ScanStatusMessage(scan_id=None, status="open", metadata={}, info={})