screenshot plot with 'S', 'Space' or 'F12'

1 files changed, 33 insertions, 0 deletions

diff --git a/picostream/dfplot.py b/picostream/dfplot.py
index aaf1ca7..50c1cba 100644
--- a/picostream/dfplot.py
+++ b/picostream/dfplot.py
@@ -580,6 +580,36 @@ class HDF5LivePlotter(QMainWindow):
                 end_time = start_time + duration
                 return np.linspace(start_time, end_time, n_samples)
 
+    def save_screenshot(self) -> None:
+        """Save a screenshot of the current plot."""
+        try:
+            timestamp = time.strftime("%Y%m%d_%H%M%S")
+            filename = f"plot_screenshot_{timestamp}.png"
+
+            # Export the plot widget as an image
+            exporter = pg.exporters.ImageExporter(self.plot_widget.plotItem)
+            exporter.export(filename)
+
+            logger.info(f"Screenshot saved: {filename}")
+
+            # Temporarily stop updates to show the message for 2 seconds
+            self.timer.stop()
+            self.acq_status_label.setText(
+                f'<span style="color: blue">Screenshot saved: {filename}</span>'
+            )
+
+            def resume_updates():
+                self.update_from_file()  # Update once immediately
+                self.timer.start(self.update_interval_ms)
+
+            QTimer.singleShot(2000, resume_updates)
+
+        except Exception as e:
+            logger.error(f"Failed to save screenshot: {e}")
+            # If we failed, restart timer if it was already running
+            if not self.timer.isActive():
+                self.timer.start(self.update_interval_ms)
+
     def closeEvent(self, event: QCloseEvent) -> None:
         """Handles the window close event.
 
@@ -595,6 +625,9 @@ class HDF5LivePlotter(QMainWindow):
         if event.key() == Qt.Key_Q:
             logger.info("'Q' key pressed. Closing application.")
             self.close()
+        elif event.key() in (Qt.Key_S, Qt.Key_Space, Qt.Key_F12):
+            logger.info("Screenshot key pressed. Saving screenshot.")
+            self.save_screenshot()
         else:
             super().keyPressEvent(event)