feat: valw window

src/main.rs

@@ -1,6 +1,3 @@
-// Window-mode helpers land before main uses them; Task 2 removes this.
-#![allow(dead_code)]
-
 mod capture;
 mod config;
 mod detach;
@@ -20,6 +17,7 @@ mod selection;
 mod stack;
 mod thumbnail;
 mod wayland;
+mod window;
 
 use std::path::PathBuf;
 use std::process::ExitCode;
@@ -59,6 +57,11 @@ enum Command {
         #[command(flatten)]
         common: Common,
     },
+    /// Click a window to capture it (Cmd+Shift+4, then Space).
+    Window {
+        #[command(flatten)]
+        common: Common,
+    },
     /// Report what the compositor and system support.
     Doctor,
     /// Internal: the process that shows preview thumbnails.
@@ -91,9 +94,14 @@ fn parse_delay(s: &str) -> Result<Duration, String> {
         .map_err(|_| format!("must be a non-negative number of seconds: {s}"))
 }
 
+/// (image, output name for the file name) pairs, which one goes to the
+/// clipboard, and the output the preview appears on.
+type Shots = (Vec<(RgbaImage, Option<String>)>, usize, String);
+
 enum Mode {
     Screen { all: bool },
     Region,
+    Window,
 }
 
 fn main() -> ExitCode {
@@ -126,6 +134,7 @@ fn run(cli: Cli) -> Result<()> {
         Command::PreviewHost => host::run(),
         Command::Screen { all, common } => capture(Mode::Screen { all }, common),
         Command::Region { common } => capture(Mode::Region, common),
+        Command::Window { common } => capture(Mode::Window, common),
     }
 }
 
@@ -145,9 +154,7 @@ fn capture(mode: Mode, common: Common) -> Result<()> {
     // Thumbnails from earlier captures must not end up in this one.
     let mut hide = ipc::HideGuard::new(&ipc::socket_path());
 
-    // (image, output name for the file name), which one goes to the
-    // clipboard, and the output the preview appears on.
-    let (shots, clip, source): (Vec<(RgbaImage, Option<String>)>, usize, String) = match mode {
+    let (shots, clip, source): Shots = match mode {
         Mode::Screen { all } => {
             let focused = focused_output(&outputs);
             let source = outputs[focused].geom.name.clone();
@@ -169,6 +176,7 @@ fn capture(mode: Mode, common: Common) -> Result<()> {
             let source = frames[i].output.name.clone();
             (vec![(frames[i].to_rgba(r), None)], 0, source)
         }
+        Mode::Window => window_shot(&outputs, cursor)?,
     };
     drop(wl);
 
@@ -195,6 +203,14 @@ fn capture(mode: Mode, common: Common) -> Result<()> {
     Ok(())
 }
 
+/// A window shot, shaped like the other modes' results. The preview goes to
+/// the window's output, or niri's focused one if that is unknown.
+fn window_shot(outputs: &[wayland::Output], cursor: bool) -> Result<Shots> {
+    let (image, output) = window::capture(cursor)?;
+    let source = output.unwrap_or_else(|| outputs[focused_output(outputs)].geom.name.clone());
+    Ok((vec![(image, None)], 0, source))
+}
+
 /// Index of niri's focused output, or 0 if niri can't tell us.
 fn focused_output(outputs: &[wayland::Output]) -> usize {
     let name = match niri::focused_output() {
@@ -284,6 +300,8 @@ mod tests {
         assert!(!parses(&["region", "--clipboard-only", "-o", "a.png"]));
         assert!(parses(&["region", "--no-preview"]));
         assert!(parses(&["screen", "--all", "--no-preview"]));
+        assert!(parses(&["window", "--cursor", "--delay", "2"]));
+        assert!(!parses(&["window", "--clipboard-only", "-o", "a.png"]));
     }
 
     #[test]

src/window.rs

@@ -0,0 +1,70 @@
+//! Window mode: niri picks the window and renders it; valw reads the result.
+
+use std::path::PathBuf;
+use std::time::Duration;
+
+use anyhow::Result;
+use image::RgbaImage;
+
+use crate::error::Cancelled;
+use crate::niri;
+
+/// How long niri gets to render and write the window.
+const TIMEOUT: Duration = Duration::from_secs(5);
+
+/// Lets the user click a window. Returns its pixels and the output it is
+/// on, if niri knows. Esc during the pick is `Cancelled`.
+pub fn capture(cursor: bool) -> Result<(RgbaImage, Option<String>)> {
+    let window = niri::pick_window()?.ok_or(Cancelled)?;
+    tracing::info!("picked window {} ({:?})", window.id, window.app_id);
+    let tmp = TempFile(temp_path());
+    niri::screenshot_window(window.id, &tmp.0, cursor, TIMEOUT)?;
+    let image = crate::thumbnail::load(&tmp.0)?;
+    let workspaces = niri::workspaces().unwrap_or_else(|e| {
+        tracing::warn!("no workspace list: {e:#}");
+        Vec::new()
+    });
+    Ok((image, niri::output_of(window.workspace_id, &workspaces)))
+}
+
+/// Where niri writes the window: next to the capture lock, one per process.
+fn temp_path() -> PathBuf {
+    crate::lock::default_path().with_file_name(format!("valw-window-{}.png", std::process::id()))
+}
+
+/// Removes the temporary screenshot however `capture` ends.
+struct TempFile(PathBuf);
+
+impl Drop for TempFile {
+    fn drop(&mut self) {
+        let _ = std::fs::remove_file(&self.0);
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn temp_file_is_removed_on_drop() {
+        let dir = std::env::temp_dir().join(format!("valw-window-test-{}", std::process::id()));
+        std::fs::create_dir_all(&dir).unwrap();
+        let path = dir.join("shot.png");
+        std::fs::write(&path, b"png").unwrap();
+        drop(TempFile(path.clone()));
+        assert!(!path.exists());
+        // A file niri never wrote is fine too.
+        drop(TempFile(dir.join("missing.png")));
+        std::fs::remove_dir(&dir).unwrap();
+    }
+
+    #[test]
+    fn temp_path_is_absolute_and_per_process() {
+        let p = temp_path();
+        assert!(p.is_absolute(), "niri rejects relative paths: {p:?}");
+        assert!(
+            p.to_string_lossy()
+                .contains(&std::process::id().to_string())
+        );
+    }
+}