Add --path argument for niri msg screenshot* commands (#2126)

* Check for empty screenshot parent before creating

Avoids a warning.

* Add --path argument for niri msg screenshot* commands

* fix

---------

Co-authored-by: Ivan Molodetskikh <yalterz@gmail.com>

6 files changed, 128 insertions, 52 deletions

diff --git a/src/input/mod.rs b/src/input/mod.rs
index 876295a2..23ea4b29 100644
--- a/src/input/mod.rs
+++ b/src/input/mod.rs
@@ -605,14 +605,17 @@ impl State {
                     self.niri.do_screen_transition(renderer, delay_ms);
                 });
             }
-            Action::ScreenshotScreen(write_to_disk, show_pointer) => {
+            Action::ScreenshotScreen(write_to_disk, show_pointer, path) => {
                 let active = self.niri.layout.active_output().cloned();
                 if let Some(active) = active {
                     self.backend.with_primary_renderer(|renderer| {
-                        if let Err(err) =
-                            self.niri
-                                .screenshot(renderer, &active, write_to_disk, show_pointer)
-                        {
+                        if let Err(err) = self.niri.screenshot(
+                            renderer,
+                            &active,
+                            write_to_disk,
+                            show_pointer,
+                            path,
+                        ) {
                             warn!("error taking screenshot: {err:?}");
                         }
                     });
@@ -636,32 +639,42 @@ impl State {
                 self.niri.screenshot_ui.toggle_pointer();
                 self.niri.queue_redraw_all();
             }
-            Action::Screenshot(show_cursor) => {
-                self.open_screenshot_ui(show_cursor);
+            Action::Screenshot(show_cursor, path) => {
+                self.open_screenshot_ui(show_cursor, path);
             }
-            Action::ScreenshotWindow(write_to_disk) => {
+            Action::ScreenshotWindow(write_to_disk, path) => {
                 let focus = self.niri.layout.focus_with_output();
                 if let Some((mapped, output)) = focus {
                     self.backend.with_primary_renderer(|renderer| {
-                        if let Err(err) =
-                            self.niri
-                                .screenshot_window(renderer, output, mapped, write_to_disk)
-                        {
+                        if let Err(err) = self.niri.screenshot_window(
+                            renderer,
+                            output,
+                            mapped,
+                            write_to_disk,
+                            path,
+                        ) {
                             warn!("error taking screenshot: {err:?}");
                         }
                     });
                 }
             }
-            Action::ScreenshotWindowById { id, write_to_disk } => {
+            Action::ScreenshotWindowById {
+                id,
+                write_to_disk,
+                path,
+            } => {
                 let mut windows = self.niri.layout.windows();
                 let window = windows.find(|(_, m)| m.id().get() == id);
                 if let Some((Some(monitor), mapped)) = window {
                     let output = monitor.output();
                     self.backend.with_primary_renderer(|renderer| {
-                        if let Err(err) =
-                            self.niri
-                                .screenshot_window(renderer, output, mapped, write_to_disk)
-                        {
+                        if let Err(err) = self.niri.screenshot_window(
+                            renderer,
+                            output,
+                            mapped,
+                            write_to_disk,
+                            path,
+                        ) {
                             warn!("error taking screenshot: {err:?}");
                         }
                     });
diff --git a/src/ipc/client.rs b/src/ipc/client.rs
index 094bb636..4c356355 100644
--- a/src/ipc/client.rs
+++ b/src/ipc/client.rs
@@ -1,20 +1,34 @@
 use std::io::ErrorKind;
 use std::iter::Peekable;
-use std::slice;
+use std::path::Path;
+use std::{env, slice};
 
 use anyhow::{anyhow, bail, Context};
 use niri_config::OutputName;
 use niri_ipc::socket::Socket;
 use niri_ipc::{
-    Event, KeyboardLayouts, LogicalOutput, Mode, Output, OutputConfigChanged, Overview, Request,
-    Response, Transform, Window, WindowLayout,
+    Action, Event, KeyboardLayouts, LogicalOutput, Mode, Output, OutputConfigChanged, Overview,
+    Request, Response, Transform, Window, WindowLayout,
 };
 use serde_json::json;
 
 use crate::cli::Msg;
 use crate::utils::version;
 
-pub fn handle_msg(msg: Msg, json: bool) -> anyhow::Result<()> {
+pub fn handle_msg(mut msg: Msg, json: bool) -> anyhow::Result<()> {
+    // For actions taking paths, prepend the niri CLI's working directory.
+    if let Msg::Action {
+        action:
+            Action::Screenshot { path, .. }
+            | Action::ScreenshotScreen { path, .. }
+            | Action::ScreenshotWindow { path, .. },
+    } = &mut msg
+    {
+        if let Some(path) = path {
+            ensure_absolute_path(path).context("error making the path absolute")?;
+        }
+    }
+
     let request = match &msg {
         Msg::Version => Request::Version,
         Msg::Outputs => Request::Outputs,
@@ -668,6 +682,19 @@ fn fmt_rounded(x: f64) -> String {
     }
 }
 
+fn ensure_absolute_path(path: &mut String) -> anyhow::Result<()> {
+    let p = Path::new(path);
+    if p.is_relative() {
+        let mut cwd = env::current_dir().context("error getting current working directory")?;
+        cwd.push(p);
+        match cwd.into_os_string().into_string() {
+            Ok(absolute) => *path = absolute,
+            Err(cwd) => bail!("couldn't convert absolute path to string: {cwd:?}"),
+        }
+    }
+    Ok(())
+}
+
 #[cfg(test)]
 mod tests {
     use insta::assert_snapshot;
diff --git a/src/ipc/server.rs b/src/ipc/server.rs
index 051bccab..354023b7 100644
--- a/src/ipc/server.rs
+++ b/src/ipc/server.rs
@@ -2,7 +2,7 @@ use std::cell::RefCell;
 use std::collections::HashSet;
 use std::ffi::OsStr;
 use std::os::unix::net::{UnixListener, UnixStream};
-use std::path::PathBuf;
+use std::path::{Path, PathBuf};
 use std::rc::Rc;
 use std::sync::{Arc, Mutex};
 use std::{env, io, process};
@@ -17,8 +17,8 @@ use futures_util::{select_biased, AsyncBufReadExt, AsyncWrite, AsyncWriteExt, Fu
 use niri_config::OutputName;
 use niri_ipc::state::{EventStreamState, EventStreamStatePart as _};
 use niri_ipc::{
-    Event, KeyboardLayouts, OutputConfigChanged, Overview, Reply, Request, Response, WindowLayout,
-    Workspace,
+    Action, Event, KeyboardLayouts, OutputConfigChanged, Overview, Reply, Request, Response,
+    WindowLayout, Workspace,
 };
 use smithay::desktop::layer_map_for_output;
 use smithay::input::pointer::{
@@ -379,6 +379,8 @@ async fn process(ctx: &ClientCtx, request: Request) -> Reply {
             Response::PickedColor(color)
         }
         Request::Action(action) => {
+            validate_action(&action)?;
+
             let (tx, rx) = async_channel::bounded(1);
 
             let action = niri_config::Action::from(action);
@@ -451,6 +453,23 @@ async fn process(ctx: &ClientCtx, request: Request) -> Reply {
     Ok(response)
 }
 
+fn validate_action(action: &Action) -> Result<(), String> {
+    if let Action::Screenshot { path, .. }
+    | Action::ScreenshotScreen { path, .. }
+    | Action::ScreenshotWindow { path, .. } = action
+    {
+        if let Some(path) = path {
+            // Relative paths are resolved against the niri compositor's working directory, which
+            // is almost certainly not what you want.
+            if !Path::new(path).is_absolute() {
+                return Err(format!("path must be absolute: {path}"));
+            }
+        }
+    }
+
+    Ok(())
+}
+
 async fn handle_event_stream_client(client: EventStreamClient) -> anyhow::Result<()> {
     let EventStreamClient {
         events,
diff --git a/src/niri.rs b/src/niri.rs
index 3ced8a90..0a7894b5 100644
--- a/src/niri.rs
+++ b/src/niri.rs
@@ -1835,7 +1835,7 @@ impl State {
         self.niri.output_management_state.notify_changes(new_config);
     }
 
-    pub fn open_screenshot_ui(&mut self, show_pointer: bool) {
+    pub fn open_screenshot_ui(&mut self, show_pointer: bool, path: Option<String>) {
         if self.niri.is_locked() || self.niri.screenshot_ui.is_open() {
             return;
         }
@@ -1876,7 +1876,7 @@ impl State {
         self.backend.with_primary_renderer(|renderer| {
             self.niri
                 .screenshot_ui
-                .open(renderer, screenshots, default_output, show_pointer)
+                .open(renderer, screenshots, default_output, show_pointer, path)
         });
 
         self.niri
@@ -1902,14 +1902,15 @@ impl State {
     }
 
     pub fn confirm_screenshot(&mut self, write_to_disk: bool) {
-        if !self.niri.screenshot_ui.is_open() {
+        let ScreenshotUi::Open { path, .. } = &mut self.niri.screenshot_ui else {
             return;
-        }
+        };
+        let path = path.take();
 
         self.backend.with_primary_renderer(|renderer| {
             match self.niri.screenshot_ui.capture(renderer) {
                 Ok((size, pixels)) => {
-                    if let Err(err) = self.niri.save_screenshot(size, pixels, write_to_disk) {
+                    if let Err(err) = self.niri.save_screenshot(size, pixels, write_to_disk, path) {
                         warn!("error saving screenshot: {err:?}");
                     }
                 }
@@ -5533,6 +5534,7 @@ impl Niri {
         output: &Output,
         write_to_disk: bool,
         include_pointer: bool,
+        path: Option<String>,
     ) -> anyhow::Result<()> {
         let _span = tracy_client::span!("Niri::screenshot");
 
@@ -5559,7 +5561,7 @@ impl Niri {
             elements,
         )?;
 
-        self.save_screenshot(size, pixels, write_to_disk)
+        self.save_screenshot(size, pixels, write_to_disk, path)
             .context("error saving screenshot")
     }
 
@@ -5569,6 +5571,7 @@ impl Niri {
         output: &Output,
         mapped: &Mapped,
         write_to_disk: bool,
+        path: Option<String>,
     ) -> anyhow::Result<()> {
         let _span = tracy_client::span!("Niri::screenshot_window");
 
@@ -5600,7 +5603,7 @@ impl Niri {
             elements,
         )?;
 
-        self.save_screenshot(geo.size, pixels, write_to_disk)
+        self.save_screenshot(geo.size, pixels, write_to_disk, path)
             .context("error saving screenshot")
     }
 
@@ -5609,14 +5612,20 @@ impl Niri {
         size: Size<i32, Physical>,
         pixels: Vec<u8>,
         write_to_disk: bool,
+        path_arg: Option<String>,
     ) -> anyhow::Result<()> {
         let path = write_to_disk
-            .then(|| match make_screenshot_path(&self.config.borrow()) {
-                Ok(path) => path,
-                Err(err) => {
-                    warn!("error making screenshot path: {err:?}");
-                    None
-                }
+            .then(|| {
+                // When given an explicit path, don't try to strftime it or create parents.
+                path_arg.map(|p| (PathBuf::from(p), false)).or_else(|| {
+                    match make_screenshot_path(&self.config.borrow()) {
+                        Ok(path) => path.map(|p| (p, true)),
+                        Err(err) => {
+                            warn!("error making screenshot path: {err:?}");
+                            None
+                        }
+                    }
+                })
             })
             .flatten();
 
@@ -5652,13 +5661,18 @@ impl Niri {
 
             let mut image_path = None;
 
-            if let Some(path) = path {
+            if let Some((path, create_parent)) = path {
                 debug!("saving screenshot to {path:?}");
 
-                if let Some(parent) = path.parent() {
-                    if let Err(err) = std::fs::create_dir(parent) {
-                        if err.kind() != std::io::ErrorKind::AlreadyExists {
-                            warn!("error creating screenshot directory: {err:?}");
+                if create_parent {
+                    if let Some(parent) = path.parent() {
+                        // Relative paths with one component, i.e. "test.png", have Some("") parent.
+                        if !parent.as_os_str().is_empty() {
+                            if let Err(err) = std::fs::create_dir(parent) {
+                                if err.kind() != std::io::ErrorKind::AlreadyExists {
+                                    warn!("error creating screenshot directory: {err:?}");
+                                }
+                            }
                         }
                     }
                 }
diff --git a/src/ui/hotkey_overlay.rs b/src/ui/hotkey_overlay.rs
index 67997fc8..541affb8 100644
--- a/src/ui/hotkey_overlay.rs
+++ b/src/ui/hotkey_overlay.rs
@@ -263,7 +263,7 @@ fn collect_actions(config: &Config) -> Vec<&Action> {
     // Screenshot is not as important, can omit if not bound.
     if let Some(bind) = binds
         .iter()
-        .find(|bind| matches!(bind.action, Action::Screenshot(_)))
+        .find(|bind| matches!(bind.action, Action::Screenshot(_, _)))
     {
         actions.push(&bind.action);
     }
@@ -479,7 +479,7 @@ fn action_name(action: &Action) -> String {
             String::from("Switch Focus Between Floating and Tiling")
         }
         Action::ToggleOverview => String::from("Open the Overview"),
-        Action::Screenshot(_) => String::from("Take a Screenshot"),
+        Action::Screenshot(_, _) => String::from("Take a Screenshot"),
         Action::Spawn(args) => format!(
             "Spawn <span face='monospace' bgcolor='#000000'>{}</span>",
             args.first().unwrap_or(&String::new())
@@ -620,7 +620,7 @@ mod tests {
     #[test]
     fn test_format_bind() {
         // Not bound.
-        assert_snapshot!(check("", Action::Screenshot(true)), @" (not bound) : Take a Screenshot");
+        assert_snapshot!(check("", Action::Screenshot(true, None)), @" (not bound) : Take a Screenshot");
 
         // Bound with a default title.
         assert_snapshot!(
@@ -628,7 +628,7 @@ mod tests {
                 r#"binds {
                     Mod+P { screenshot; }
                 }"#,
-                Action::Screenshot(true),
+                Action::Screenshot(true, None),
             ),
             @" Super + P : Take a Screenshot"
         );
@@ -639,7 +639,7 @@ mod tests {
                 r#"binds {
                     Mod+P hotkey-overlay-title="Hello" { screenshot; }
                 }"#,
-                Action::Screenshot(true),
+                Action::Screenshot(true, None),
             ),
             @" Super + P : Hello"
         );
@@ -651,7 +651,7 @@ mod tests {
                     Mod+P { screenshot; }
                     Print { screenshot; }
                 }"#,
-                Action::Screenshot(true),
+                Action::Screenshot(true, None),
             ),
             @" Super + P : Take a Screenshot"
         );
@@ -663,7 +663,7 @@ mod tests {
                     Mod+P { screenshot; }
                     Print hotkey-overlay-title="My Cool Bind" { screenshot; }
                 }"#,
-                Action::Screenshot(true),
+                Action::Screenshot(true, None),
             ),
             @" PrtSc : My Cool Bind"
         );
@@ -675,7 +675,7 @@ mod tests {
                     Mod+P hotkey-overlay-title="First" { screenshot; }
                     Print hotkey-overlay-title="My Cool Bind" { screenshot; }
                 }"#,
-                Action::Screenshot(true),
+                Action::Screenshot(true, None),
             ),
             @" Super + P : First"
         );
@@ -687,7 +687,7 @@ mod tests {
                     Mod+P { screenshot; }
                     Print hotkey-overlay-title=null { screenshot; }
                 }"#,
-                Action::Screenshot(true),
+                Action::Screenshot(true, None),
             ),
             @"None"
         );
@@ -699,7 +699,7 @@ mod tests {
                     Mod+P hotkey-overlay-title="Hello" { screenshot; }
                     Print hotkey-overlay-title=null { screenshot; }
                 }"#,
-                Action::Screenshot(true),
+                Action::Screenshot(true, None),
             ),
             @" Super + P : Hello"
         );
diff --git a/src/ui/screenshot_ui.rs b/src/ui/screenshot_ui.rs
index 3200cae4..7c03d36e 100644
--- a/src/ui/screenshot_ui.rs
+++ b/src/ui/screenshot_ui.rs
@@ -63,6 +63,7 @@ pub enum ScreenshotUi {
         open_anim: Animation,
         clock: Clock,
         config: Rc<RefCell<Config>>,
+        path: Option<String>,
     },
 }
 
@@ -141,6 +142,7 @@ impl ScreenshotUi {
         screenshots: HashMap<Output, [OutputScreenshot; 3]>,
         default_output: Output,
         show_pointer: bool,
+        path: Option<String>,
     ) -> bool {
         if screenshots.is_empty() {
             return false;
@@ -235,6 +237,7 @@ impl ScreenshotUi {
             open_anim,
             clock: clock.clone(),
             config: config.clone(),
+            path,
         };
 
         self.update_buffers();