diff --git a/src/main.rs b/src/main.rs index 5d367cc..0707a9b 100644 --- a/src/main.rs +++ b/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 { .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)>, 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)>, 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 { + 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] diff --git a/src/window.rs b/src/window.rs new file mode 100644 index 0000000..e65a504 --- /dev/null +++ b/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)> { + 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()) + ); + } +}