feat: add scan_id parameter and update_with_scan_history method to Heatmap class

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":

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.

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={})