Allow attaching custom user data to a screenshot command (#5416)

This lets users trigger a screenshot from anywhere, and then when they get back the results they have some context about what part of their code triggered the screenshot.
emilk · Dec 3, 2024 · a9c76ba · a9c76ba
1 parent 6a1131f
commit a9c76ba
Show file tree

Hide file tree

Showing 9 changed files with 129 additions and 23 deletions.
diff --git a/crates/eframe/src/native/glow_integration.rs b/crates/eframe/src/native/glow_integration.rs
@@ -661,13 +661,14 @@ impl<'app> GlowWinitRunning<'app> {
         {
             for action in viewport.actions_requested.drain() {
                 match action {
-                    ActionRequested::Screenshot => {
+                    ActionRequested::Screenshot(user_data) => {
                         let screenshot = painter.read_screen_rgba(screen_size_in_pixels);
                         egui_winit
                             .egui_input_mut()
                             .events
                             .push(egui::Event::Screenshot {
                                 viewport_id,
+                                user_data,
                                 image: screenshot.into(),
                             });
                     }

diff --git a/crates/eframe/src/native/wgpu_integration.rs b/crates/eframe/src/native/wgpu_integration.rs
@@ -643,10 +643,16 @@ impl<'app> WgpuWinitRunning<'app> {
 
         let clipped_primitives = egui_ctx.tessellate(shapes, pixels_per_point);
 
-        let screenshot_requested = viewport
-            .actions_requested
-            .take(&ActionRequested::Screenshot)
-            .is_some();
+        let mut screenshot_commands = vec![];
+        viewport.actions_requested.retain(|cmd| {
+            if let ActionRequested::Screenshot(info) = cmd {
+                screenshot_commands.push(info.clone());
+                false
+            } else {
+                true
+            }
+        });
+        let screenshot_requested = !screenshot_commands.is_empty();
         let (vsync_secs, screenshot) = painter.paint_and_update_textures(
             viewport_id,
             pixels_per_point,
@@ -655,19 +661,32 @@ impl<'app> WgpuWinitRunning<'app> {
             &textures_delta,
             screenshot_requested,
         );
-        if let Some(screenshot) = screenshot {
-            egui_winit
-                .egui_input_mut()
-                .events
-                .push(egui::Event::Screenshot {
-                    viewport_id,
-                    image: screenshot.into(),
-                });
+        match (screenshot_requested, screenshot) {
+            (false, None) => {}
+            (true, Some(screenshot)) => {
+                let screenshot = Arc::new(screenshot);
+                for user_data in screenshot_commands {
+                    egui_winit
+                        .egui_input_mut()
+                        .events
+                        .push(egui::Event::Screenshot {
+                            viewport_id,
+                            user_data,
+                            image: screenshot.clone(),
+                        });
+                }
+            }
+            (true, None) => {
+                log::error!("Bug in egui_wgpu: screenshot requested, but no screenshot was taken");
+            }
+            (false, Some(_)) => {
+                log::warn!("Bug in egui_wgpu: Got screenshot without requesting it");
+            }
         }
 
         for action in viewport.actions_requested.drain() {
             match action {
-                ActionRequested::Screenshot => {
+                ActionRequested::Screenshot { .. } => {
                     // already handled above
                 }
                 ActionRequested::Cut => {

diff --git a/crates/egui-winit/src/lib.rs b/crates/egui-winit/src/lib.rs
@@ -1301,7 +1301,7 @@ fn translate_cursor(cursor_icon: egui::CursorIcon) -> Option<winit::window::Curs
 // ---------------------------------------------------------------------------
 #[derive(PartialEq, Eq, Hash, Debug)]
 pub enum ActionRequested {
-    Screenshot,
+    Screenshot(egui::UserData),
     Cut,
     Copy,
     Paste,
@@ -1516,8 +1516,8 @@ fn process_viewport_command(
                 log::warn!("{command:?}: {err}");
             }
         }
-        ViewportCommand::Screenshot => {
-            actions_requested.insert(ActionRequested::Screenshot);
+        ViewportCommand::Screenshot(user_data) => {
+            actions_requested.insert(ActionRequested::Screenshot(user_data));
         }
         ViewportCommand::RequestCut => {
             actions_requested.insert(ActionRequested::Cut);

diff --git a/crates/egui/src/data/input.rs b/crates/egui/src/data/input.rs
@@ -529,6 +529,10 @@ pub enum Event {
     /// The reply of a screenshot requested with [`crate::ViewportCommand::Screenshot`].
     Screenshot {
         viewport_id: crate::ViewportId,
+
+        /// Whatever was passed to [`crate::ViewportCommand::Screenshot`].
+        user_data: crate::UserData,
+
         image: std::sync::Arc<ColorImage>,
     },
 }

diff --git a/crates/egui/src/data/mod.rs b/crates/egui/src/data/mod.rs
@@ -3,5 +3,7 @@
 pub mod input;
 mod key;
 pub mod output;
+mod user_data;
 
 pub use key::Key;
+pub use user_data::UserData;
diff --git a/crates/egui/src/data/user_data.rs b/crates/egui/src/data/user_data.rs
@@ -0,0 +1,74 @@
+use std::{any::Any, sync::Arc};
+
+/// A wrapper around `dyn Any`, used for passing custom user data
+/// to [`crate::ViewportCommand::Screenshot`].
+#[derive(Clone, Debug, Default)]
+pub struct UserData {
+    /// A user value given to the screenshot command,
+    /// that will be returned in [`crate::Event::Screenshot`].
+    pub data: Option<Arc<dyn Any + Send + Sync>>,
+}
+
+impl UserData {
+    /// You can also use [`Self::default`].
+    pub fn new(user_info: impl Any + Send + Sync) -> Self {
+        Self {
+            data: Some(Arc::new(user_info)),
+        }
+    }
+}
+
+impl PartialEq for UserData {
+    fn eq(&self, other: &Self) -> bool {
+        match (&self.data, &other.data) {
+            (Some(a), Some(b)) => Arc::ptr_eq(a, b),
+            (None, None) => true,
+            _ => false,
+        }
+    }
+}
+
+impl Eq for UserData {}
+
+impl std::hash::Hash for UserData {
+    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+        self.data.as_ref().map(Arc::as_ptr).hash(state);
+    }
+}
+
+#[cfg(feature = "serde")]
+impl serde::Serialize for UserData {
+    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: serde::Serializer,
+    {
+        serializer.serialize_none() // can't serialize an `Any`
+    }
+}
+
+#[cfg(feature = "serde")]
+impl<'de> serde::Deserialize<'de> for UserData {
+    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+    where
+        D: serde::Deserializer<'de>,
+    {
+        struct UserDataVisitor;
+
+        impl<'de> serde::de::Visitor<'de> for UserDataVisitor {
+            type Value = UserData;
+
+            fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+                formatter.write_str("a None value")
+            }
+
+            fn visit_none<E>(self) -> Result<UserData, E>
+            where
+                E: serde::de::Error,
+            {
+                Ok(UserData::default())
+            }
+        }
+
+        deserializer.deserialize_option(UserDataVisitor)
+    }
+}
diff --git a/crates/egui/src/lib.rs b/crates/egui/src/lib.rs
@@ -471,7 +471,7 @@ pub use self::{
         output::{
             self, CursorIcon, FullOutput, OpenUrl, PlatformOutput, UserAttentionType, WidgetInfo,
         },
-        Key,
+        Key, UserData,
     },
     drag_and_drop::DragAndDrop,
     epaint::text::TextWrapMode,

diff --git a/crates/egui/src/viewport.rs b/crates/egui/src/viewport.rs
@@ -1058,8 +1058,8 @@ pub enum ViewportCommand {
 
     /// Take a screenshot.
     ///
-    /// The results are returned in `crate::Event::Screenshot`.
-    Screenshot,
+    /// The results are returned in [`crate::Event::Screenshot`].
+    Screenshot(crate::UserData),
 
     /// Request cut of the current selection
     ///
@@ -1100,6 +1100,8 @@ impl ViewportCommand {
     }
 }
 
+// ----------------------------------------------------------------------------
+
 /// Describes a viewport, i.e. a native window.
 ///
 /// This is returned by [`crate::Context::run`] on each frame, and should be applied

diff --git a/examples/screenshot/src/main.rs b/examples/screenshot/src/main.rs
@@ -45,7 +45,7 @@ impl eframe::App for MyApp {
 
                 if ui.button("save to 'top_left.png'").clicked() {
                     self.save_to_file = true;
-                    ctx.send_viewport_cmd(egui::ViewportCommand::Screenshot);
+                    ctx.send_viewport_cmd(egui::ViewportCommand::Screenshot(Default::default()));
                 }
 
                 ui.with_layout(egui::Layout::top_down(egui::Align::RIGHT), |ui| {
@@ -58,9 +58,13 @@ impl eframe::App for MyApp {
                         } else {
                             ctx.set_theme(egui::Theme::Light);
                         };
-                        ctx.send_viewport_cmd(egui::ViewportCommand::Screenshot);
+                        ctx.send_viewport_cmd(
+                            egui::ViewportCommand::Screenshot(Default::default()),
+                        );
                     } else if ui.button("take screenshot!").clicked() {
-                        ctx.send_viewport_cmd(egui::ViewportCommand::Screenshot);
+                        ctx.send_viewport_cmd(
+                            egui::ViewportCommand::Screenshot(Default::default()),
+                        );
                     }
                 });
             });