Add screenshot-window action

5 files changed, 73 insertions, 0 deletions

diff --git a/resources/default-config.kdl b/resources/default-config.kdl
index aab1d77d..0c27a45d 100644
--- a/resources/default-config.kdl
+++ b/resources/default-config.kdl
@@ -208,6 +208,8 @@ binds {
     Mod+Equal { set-column-width "+10%"; }
 
     Print { screenshot; }
+    Alt+Print { screenshot-window; }
+
     Mod+Shift+E { quit; }
     Mod+Shift+P { power-off-monitors; }
 
diff --git a/src/config.rs b/src/config.rs
index 4ab9c36e..654b0557 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -239,6 +239,7 @@ pub enum Action {
     ToggleDebugTint,
     Spawn(#[knuffel(arguments)] Vec<String>),
     Screenshot,
+    ScreenshotWindow,
     CloseWindow,
     FullscreenWindow,
     FocusColumnLeft,
diff --git a/src/input.rs b/src/input.rs
index c1e335dd..9f59c06e 100644
--- a/src/input.rs
+++ b/src/input.rs
@@ -167,6 +167,18 @@ impl State {
                                 }
                             }
                         }
+                        Action::ScreenshotWindow => {
+                            let active = self.niri.layout.active_window();
+                            if let Some((window, output)) = active {
+                                if let Some(renderer) = self.backend.renderer() {
+                                    if let Err(err) =
+                                        self.niri.screenshot_window(renderer, &output, &window)
+                                    {
+                                        warn!("error taking screenshot: {err:?}");
+                                    }
+                                }
+                            }
+                        }
                         Action::CloseWindow => {
                             if let Some(window) = self.niri.layout.focus() {
                                 window.toplevel().send_close();
diff --git a/src/layout.rs b/src/layout.rs
index 003bb1c4..85d90bfb 100644
--- a/src/layout.rs
+++ b/src/layout.rs
@@ -751,6 +751,30 @@ impl<W: LayoutElement> Layout<W> {
         Some(&mon.workspaces[mon.active_workspace_idx])
     }
 
+    pub fn active_window(&self) -> Option<(W, Output)> {
+        let MonitorSet::Normal {
+            monitors,
+            active_monitor_idx,
+            ..
+        } = &self.monitor_set
+        else {
+            return None;
+        };
+
+        let mon = &monitors[*active_monitor_idx];
+        let ws = &mon.workspaces[mon.active_workspace_idx];
+
+        if ws.columns.is_empty() {
+            return None;
+        }
+
+        let col = &ws.columns[ws.active_column_idx];
+        Some((
+            col.windows[col.active_window_idx].clone(),
+            mon.output.clone(),
+        ))
+    }
+
     pub fn workspace_for_output(&self, output: &Output) -> Option<&Workspace<W>> {
         let MonitorSet::Normal { monitors, .. } = &self.monitor_set else {
             return None;
diff --git a/src/niri.rs b/src/niri.rs
index 3fbb35f0..8d437de1 100644
--- a/src/niri.rs
+++ b/src/niri.rs
@@ -1510,6 +1510,40 @@ impl Niri {
         let elements = self.render(renderer, output, true);
         let pixels = render_to_vec(renderer, size, scale, &elements)?;
 
+        self.save_screenshot(size, pixels)
+            .context("error saving screenshot")
+    }
+
+    pub fn screenshot_window(
+        &mut self,
+        renderer: &mut GlesRenderer,
+        output: &Output,
+        window: &Window,
+    ) -> anyhow::Result<()> {
+        let _span = tracy_client::span!("Niri::screenshot_window");
+
+        let scale = Scale::from(output.current_scale().fractional_scale());
+        let bbox = window.bbox_with_popups();
+        let size = bbox.size.to_physical_precise_ceil(scale);
+        let buf_pos = Point::from((0, 0)) - bbox.loc;
+        // FIXME: pointer.
+        let elements = window.render_elements::<WaylandSurfaceRenderElement<GlesRenderer>>(
+            renderer,
+            buf_pos.to_physical_precise_ceil(scale),
+            scale,
+            1.,
+        );
+        let pixels = render_to_vec(renderer, size, scale, &elements)?;
+
+        self.save_screenshot(size, pixels)
+            .context("error saving screenshot")
+    }
+
+    fn save_screenshot(
+        &mut self,
+        size: Size<i32, Physical>,
+        pixels: Vec<u8>,
+    ) -> anyhow::Result<()> {
         let path = make_screenshot_path().context("error making screenshot path")?;
         debug!("saving screenshot to {path:?}");