From 68a5c6ac1ce6373f43abe9703e8c515e840ab4d7 Mon Sep 17 00:00:00 2001 From: Access Date: Fri, 2 Feb 2024 19:18:37 +0000 Subject: [PATCH 1/6] Refactor: introduce FrameGuard, ScreenCapturer, and Logical/EmbeddedRegion --- libwayshot/src/dispatch.rs | 29 ++-- libwayshot/src/image_util.rs | 8 + libwayshot/src/lib.rs | 325 ++++++++++++++++++----------------- libwayshot/src/region.rs | 219 +++++++++++++++++++++++ libwayshot/src/screencopy.rs | 26 ++- wayshot/src/utils.rs | 16 +- 6 files changed, 435 insertions(+), 188 deletions(-) create mode 100644 libwayshot/src/region.rs diff --git a/libwayshot/src/dispatch.rs b/libwayshot/src/dispatch.rs index 9b760c21..170e1026 100644 --- a/libwayshot/src/dispatch.rs +++ b/libwayshot/src/dispatch.rs @@ -25,11 +25,13 @@ use crate::{ screencopy::FrameFormat, }; +#[derive(Debug)] pub struct OutputCaptureState { pub outputs: Vec, } impl Dispatch for OutputCaptureState { + #[tracing::instrument(skip(wl_registry, qh), ret, level = "trace")] fn event( state: &mut Self, wl_registry: &WlRegistry, @@ -76,6 +78,7 @@ impl Dispatch for OutputCaptureState { } impl Dispatch for OutputCaptureState { + #[tracing::instrument(skip(wl_output), ret, level = "trace")] fn event( state: &mut Self, wl_output: &WlOutput, @@ -114,6 +117,7 @@ impl Dispatch for OutputCaptureState { delegate_noop!(OutputCaptureState: ignore ZxdgOutputManagerV1); impl Dispatch for OutputCaptureState { + #[tracing::instrument(ret, level = "trace")] fn event( state: &mut Self, _: &ZxdgOutputV1, @@ -128,19 +132,20 @@ impl Dispatch for OutputCaptureState { zxdg_output_v1::Event::LogicalPosition { x, y } => { output_info.dimensions.x = x; output_info.dimensions.y = y; - tracing::debug!("Logical position event fired!"); } zxdg_output_v1::Event::LogicalSize { width, height } => { output_info.dimensions.width = width; output_info.dimensions.height = height; - tracing::debug!("Logical size event fired!"); } + zxdg_output_v1::Event::Done => {} + zxdg_output_v1::Event::Name { .. } => {} + zxdg_output_v1::Event::Description { .. } => {} _ => {} }; } } -/// State of the frame after attemting to copy it's data to a wl_buffer. +/// State of the frame after attempting to copy it's data to a wl_buffer. #[derive(Debug, Copy, Clone, PartialEq)] pub enum FrameState { /// Compositor returned a failed event on calling `frame.copy`. @@ -156,6 +161,7 @@ pub struct CaptureFrameState { } impl Dispatch for CaptureFrameState { + #[tracing::instrument(skip(frame), ret, level = "trace")] fn event( frame: &mut Self, _: &ZwlrScreencopyFrameV1, @@ -171,7 +177,6 @@ impl Dispatch for CaptureFrameState { height, stride, } => { - tracing::debug!("Received Buffer event"); if let Value(f) = format { frame.formats.push(FrameFormat { format: f, @@ -184,30 +189,20 @@ impl Dispatch for CaptureFrameState { exit(1); } } - zwlr_screencopy_frame_v1::Event::Flags { .. } => { - tracing::debug!("Received Flags event"); - } zwlr_screencopy_frame_v1::Event::Ready { .. } => { // If the frame is successfully copied, a “flags” and a “ready” events are sent. Otherwise, a “failed” event is sent. // This is useful when we call .copy on the frame object. - tracing::debug!("Received Ready event"); frame.state.replace(FrameState::Finished); } zwlr_screencopy_frame_v1::Event::Failed => { - tracing::debug!("Received Failed event"); frame.state.replace(FrameState::Failed); } - zwlr_screencopy_frame_v1::Event::Damage { .. } => { - tracing::debug!("Received Damage event"); - } - zwlr_screencopy_frame_v1::Event::LinuxDmabuf { .. } => { - tracing::debug!("Received LinuxDmaBuf event"); - } + zwlr_screencopy_frame_v1::Event::Damage { .. } => {} + zwlr_screencopy_frame_v1::Event::LinuxDmabuf { .. } => {} zwlr_screencopy_frame_v1::Event::BufferDone => { - tracing::debug!("Received bufferdone event"); frame.buffer_done.store(true, Ordering::SeqCst); } - _ => unreachable!(), + _ => {} }; } } diff --git a/libwayshot/src/image_util.rs b/libwayshot/src/image_util.rs index 92eb3967..197e8f11 100644 --- a/libwayshot/src/image_util.rs +++ b/libwayshot/src/image_util.rs @@ -7,6 +7,14 @@ pub(crate) fn rotate_image_buffer( width: u32, height: u32, ) -> DynamicImage { + // TODO Better document whether width and height are before or after the transform. + // Perhaps this should be part of a cleanup of the FrameCopy struct. + let (width, height) = match transform { + Transform::_90 | Transform::_270 | Transform::Flipped90 | Transform::Flipped270 => { + (height, width) + } + _ => (width, height), + }; let final_image = match transform { Transform::_90 => image::imageops::rotate90(&image).into(), Transform::_180 => image::imageops::rotate180(&image).into(), diff --git a/libwayshot/src/lib.rs b/libwayshot/src/lib.rs index 6c003320..c93cb335 100644 --- a/libwayshot/src/lib.rs +++ b/libwayshot/src/lib.rs @@ -8,10 +8,10 @@ mod dispatch; mod error; mod image_util; pub mod output; +pub mod region; mod screencopy; use std::{ - cmp, fs::File, os::fd::AsFd, process::exit, @@ -19,12 +19,15 @@ use std::{ thread, }; -use image::{imageops::overlay, DynamicImage}; +use image::{imageops::replace, DynamicImage}; use memmap2::MmapMut; +use region::{EmbeddedRegion, RegionCapturer}; +use screencopy::FrameGuard; +use tracing::debug; use wayland_client::{ globals::{registry_queue_init, GlobalList}, protocol::{ - wl_output::{Transform, WlOutput}, + wl_output::WlOutput, wl_shm::{self, WlShm}, }, Connection, EventQueue, @@ -41,6 +44,7 @@ use crate::{ convert::create_converter, dispatch::{CaptureFrameState, FrameState, OutputCaptureState, WayshotState}, output::OutputInfo, + region::LogicalRegion, screencopy::{create_shm_fd, FrameCopy, FrameFormat}, }; @@ -51,28 +55,6 @@ pub mod reexport { pub use wl_output::{Transform, WlOutput}; } -type Frame = (Vec, (i32, i32)); - -/// Struct to store region capture details. -#[derive(Debug, Copy, Clone, PartialEq, Eq)] -pub struct CaptureRegion { - /// X coordinate of the area to capture. - pub x_coordinate: i32, - /// y coordinate of the area to capture. - pub y_coordinate: i32, - /// Width of the capture area. - pub width: i32, - /// Height of the capture area. - pub height: i32, -} - -#[derive(Debug)] -struct IntersectingOutput { - output: WlOutput, - region: CaptureRegion, - transform: Transform, -} - /// Struct to store wayland connection and globals list. /// # Example usage /// @@ -161,7 +143,7 @@ impl WayshotConnection { tracing::error!("Compositor did not advertise any wl_output devices!"); exit(1); } - tracing::debug!("Outputs detected: {:#?}", state.outputs); + tracing::trace!("Outputs detected: {:#?}", state.outputs); self.output_infos = state.outputs; Ok(()) @@ -174,18 +156,21 @@ impl WayshotConnection { cursor_overlay: i32, output: &WlOutput, fd: T, - capture_region: Option, - ) -> Result { + capture_region: Option, + ) -> Result<(FrameFormat, FrameGuard)> { let (state, event_queue, frame, frame_format) = self.capture_output_frame_get_state(cursor_overlay, output, capture_region)?; - self.capture_output_frame_inner(state, event_queue, frame, frame_format, fd) + let frame_guard = + self.capture_output_frame_inner(state, event_queue, frame, frame_format, fd)?; + + Ok((frame_format, frame_guard)) } fn capture_output_frame_get_state( &self, cursor_overlay: i32, output: &WlOutput, - capture_region: Option, + capture_region: Option, ) -> Result<( CaptureFrameState, EventQueue, @@ -216,15 +201,15 @@ impl WayshotConnection { } }; - // Capture output. - let frame: ZwlrScreencopyFrameV1 = if let Some(region) = capture_region { + debug!("Capturing output..."); + let frame = if let Some(embedded_region) = capture_region { screencopy_manager.capture_output_region( cursor_overlay, output, - region.x_coordinate, - region.y_coordinate, - region.width, - region.height, + embedded_region.inner.x, + embedded_region.inner.y, + embedded_region.inner.width, + embedded_region.inner.height, &qh, (), ) @@ -238,7 +223,7 @@ impl WayshotConnection { event_queue.blocking_dispatch(&mut state)?; } - tracing::debug!( + tracing::trace!( "Received compositor frame buffer formats: {:#?}", state.formats ); @@ -258,7 +243,7 @@ impl WayshotConnection { ) }) .copied(); - tracing::debug!("Selected frame buffer format: {:#?}", frame_format); + tracing::trace!("Selected frame buffer format: {:#?}", frame_format); // Check if frame format exists. let frame_format = match frame_format { @@ -278,7 +263,7 @@ impl WayshotConnection { frame: ZwlrScreencopyFrameV1, frame_format: FrameFormat, fd: T, - ) -> Result { + ) -> Result { // Connecting to wayland environment. let qh = event_queue.handle(); @@ -310,9 +295,7 @@ impl WayshotConnection { return Err(Error::FramecopyFailed); } FrameState::Finished => { - buffer.destroy(); - shm_pool.destroy(); - return Ok(frame_format); + return Ok(FrameGuard { buffer, shm_pool }); } } } @@ -326,8 +309,8 @@ impl WayshotConnection { cursor_overlay: bool, output: &WlOutput, file: &File, - capture_region: Option, - ) -> Result { + capture_region: Option, + ) -> Result<(FrameFormat, FrameGuard)> { let (state, event_queue, frame, frame_format) = self.capture_output_frame_get_state(cursor_overlay as i32, output, capture_region)?; @@ -335,25 +318,28 @@ impl WayshotConnection { let frame_bytes = frame_format.stride * frame_format.height; file.set_len(frame_bytes as u64)?; - self.capture_output_frame_inner(state, event_queue, frame, frame_format, file) + let frame_guard = + self.capture_output_frame_inner(state, event_queue, frame, frame_format, file)?; + + Ok((frame_format, frame_guard)) } /// Get a FrameCopy instance with screenshot pixel data for any wl_output object. - fn capture_output_frame( + #[tracing::instrument(skip_all, fields(output = output_info.name, region = capture_region.map(|r| format!("{:}", r))))] + fn capture_frame_copy( &self, cursor_overlay: bool, - output: &WlOutput, - transform: Transform, - capture_region: Option, - ) -> Result { + output_info: &OutputInfo, + capture_region: Option, + ) -> Result<(FrameCopy, FrameGuard)> { // Create an in memory file and return it's file descriptor. let fd = create_shm_fd()?; // Create a writeable memory map backed by a mem_file. let mem_file = File::from(fd); - let frame_format = self.capture_output_frame_shm_from_file( + let (frame_format, frame_guard) = self.capture_output_frame_shm_from_file( cursor_overlay, - output, + &output_info.wl_output, &mem_file, capture_region, )?; @@ -367,64 +353,45 @@ impl WayshotConnection { tracing::error!("You can send a feature request for the above format to the mailing list for wayshot over at https://sr.ht/~shinyzenith/wayshot."); return Err(Error::NoSupportedBufferFormat); }; - Ok(FrameCopy { + let frame_copy = FrameCopy { frame_format, frame_color_type, frame_mmap, - transform, - }) + transform: output_info.transform, + position: capture_region + .map(|capture_region| { + let logical_region = capture_region.logical(); + (logical_region.inner.x as i64, logical_region.inner.y as i64) + }) + .unwrap_or_else(|| { + ( + output_info.dimensions.x as i64, + output_info.dimensions.y as i64, + ) + }), + }; + tracing::debug!("Created frame copy: {:#?}", frame_copy); + Ok((frame_copy, frame_guard)) } - fn create_frame_copy( + pub fn capture_frame_copies( &self, - capture_region: CaptureRegion, + output_capture_regions: &Vec<(OutputInfo, Option)>, cursor_overlay: bool, - ) -> Result { + ) -> Result> { let frame_copies = thread::scope(|scope| -> Result<_> { - let join_handles = self - .get_all_outputs() + let join_handles = output_capture_regions .into_iter() - .filter_map(|output| { - let x1: i32 = cmp::max(output.dimensions.x, capture_region.x_coordinate); - let y1: i32 = cmp::max(output.dimensions.y, capture_region.y_coordinate); - let x2: i32 = cmp::min( - output.dimensions.x + output.dimensions.width, - capture_region.x_coordinate + capture_region.width, - ); - let y2: i32 = cmp::min( - output.dimensions.y + output.dimensions.height, - capture_region.y_coordinate + capture_region.height, - ); - - let width = x2 - x1; - let height = y2 - y1; - - if width <= 0 || height <= 0 { - return None; - } - - let true_x = capture_region.x_coordinate - output.dimensions.x; - let true_y = capture_region.y_coordinate - output.dimensions.y; - let true_region = CaptureRegion { - x_coordinate: true_x, - y_coordinate: true_y, - width: capture_region.width, - height: capture_region.height, - }; - Some(IntersectingOutput { - output: output.wl_output.clone(), - region: true_region, - transform: output.transform, - }) - }) - .map(|intersecting_output| { + .map(|(output_info, capture_region)| { scope.spawn(move || { - self.capture_output_frame( + self.capture_frame_copy( cursor_overlay, - &intersecting_output.output, - intersecting_output.transform, - Some(intersecting_output.region), + &output_info, + capture_region.clone(), ) + .map(|(frame_copy, frame_guard)| { + (frame_copy, frame_guard, output_info.clone()) + }) }) }) .collect::>(); @@ -436,30 +403,71 @@ impl WayshotConnection { .collect::>() })?; - Ok((frame_copies, (capture_region.width, capture_region.height))) + Ok(frame_copies) } /// Take a screenshot from the specified region. - pub fn screenshot( + fn screenshot_region_capturer( &self, - capture_region: CaptureRegion, + region_capturer: RegionCapturer, cursor_overlay: bool, ) -> Result { - let (frame_copies, (width, height)) = - self.create_frame_copy(capture_region, cursor_overlay)?; + let outputs_capture_regions: &Vec<(OutputInfo, Option)> = + &match region_capturer { + RegionCapturer::Outputs(ref outputs) => outputs + .into_iter() + .map(|output_info| (output_info.clone(), None)) + .collect(), + RegionCapturer::Region(capture_region) => self + .get_all_outputs() + .into_iter() + .filter_map(|output_info| { + tracing::span!( + tracing::Level::DEBUG, + "filter_map", + output = format!( + "{name} at {region}", + name = output_info.name, + region = LogicalRegion::from(output_info), + ), + capture_region = format!("{}", capture_region), + ) + .in_scope(|| { + if let Some(relative_region) = + EmbeddedRegion::new(capture_region, output_info.into()) + { + tracing::debug!("Intersection found: {}", relative_region); + Some((output_info.clone(), Some(relative_region))) + } else { + tracing::debug!("No intersection found"); + None + } + }) + }) + .collect(), + }; + + let frames = self.capture_frame_copies(outputs_capture_regions, cursor_overlay)?; + + let capture_region: LogicalRegion = match region_capturer { + RegionCapturer::Outputs(ref outputs) => outputs.try_into()?, + RegionCapturer::Region(region) => region, + }; thread::scope(|scope| { - let rotate_join_handles = frame_copies + let rotate_join_handles = frames .into_iter() - .map(|frame_copy| { + .map(|(frame_copy, _, _)| { scope.spawn(move || { - let transform = frame_copy.transform; - let image = frame_copy.try_into()?; - Ok(image_util::rotate_image_buffer( - image, - transform, - width as u32, - height as u32, + let image = (&frame_copy).try_into()?; + Ok(( + image_util::rotate_image_buffer( + image, + frame_copy.transform, + frame_copy.frame_format.width, + frame_copy.frame_format.height, + ), + frame_copy, )) }) }) @@ -471,21 +479,38 @@ impl WayshotConnection { .flatten() .fold( None, - |possible_overlayed_image_or_error: Option>, image: Result<_>| { - if let Some(overlayed_image_or_error) = possible_overlayed_image_or_error { - if let Ok(mut overlayed_image) = overlayed_image_or_error { - if let Ok(image) = image { - overlay(&mut overlayed_image, &image, 0, 0); - Some(Ok(overlayed_image)) - } else { - Some(image) - } - } else { - Some(image) - } - } else { - Some(image) - } + |composite_image: Option>, image: Result<_>| { + // Default to a transparent image. + let composite_image = composite_image.unwrap_or_else(|| { + Ok(DynamicImage::new_rgba8( + capture_region.inner.width as u32, + capture_region.inner.height as u32, + )) + }); + + Some(|| -> Result<_> { + let mut composite_image = composite_image?; + let (image, frame_copy) = image?; + let frame_copy_region = LogicalRegion::from(&frame_copy); + let (x, y) = ( + frame_copy_region.inner.x as i64 - capture_region.inner.x as i64, + frame_copy_region.inner.y as i64 - capture_region.inner.y as i64, + ); + tracing::span!( + tracing::Level::DEBUG, + "replace", + frame_copy_region = format!("{}", frame_copy_region), + capture_region = format!("{}", capture_region), + x = x, + y = y, + ) + .in_scope(|| { + tracing::debug!("Replacing parts of the final image"); + replace(&mut composite_image, &image, x, y); + }); + + Ok(composite_image) + }()) }, ) .ok_or_else(|| { @@ -495,19 +520,23 @@ impl WayshotConnection { }) } + /// Take a screenshot from the specified region. + pub fn screenshot( + &self, + capture_region: LogicalRegion, + cursor_overlay: bool, + ) -> Result { + self.screenshot_region_capturer(RegionCapturer::Region(capture_region), cursor_overlay) + } + /// shot one ouput pub fn screenshot_single_output( &self, output_info: &OutputInfo, cursor_overlay: bool, ) -> Result { - let frame_copy = self.capture_output_frame( - cursor_overlay, - &output_info.wl_output, - output_info.transform, - None, - )?; - frame_copy.try_into() + let (frame_copy, _) = self.capture_frame_copy(cursor_overlay, output_info, None)?; + (&frame_copy).try_into() } /// Take a screenshot from all of the specified outputs. @@ -520,33 +549,7 @@ impl WayshotConnection { return Err(Error::NoOutputs); } - let x1 = outputs - .iter() - .map(|output| output.dimensions.x) - .min() - .unwrap(); - let y1 = outputs - .iter() - .map(|output| output.dimensions.y) - .min() - .unwrap(); - let x2 = outputs - .iter() - .map(|output| output.dimensions.x + output.dimensions.width) - .max() - .unwrap(); - let y2 = outputs - .iter() - .map(|output| output.dimensions.y + output.dimensions.height) - .max() - .unwrap(); - let capture_region = CaptureRegion { - x_coordinate: x1, - y_coordinate: y1, - width: x2 - x1, - height: y2 - y1, - }; - self.screenshot(capture_region, cursor_overlay) + self.screenshot_region_capturer(RegionCapturer::Outputs(outputs.clone()), cursor_overlay) } /// Take a screenshot from all accessible outputs. diff --git a/libwayshot/src/region.rs b/libwayshot/src/region.rs new file mode 100644 index 00000000..a82e0cdc --- /dev/null +++ b/libwayshot/src/region.rs @@ -0,0 +1,219 @@ +use std::cmp; + +use wayland_client::protocol::wl_output::Transform; + +use crate::error::Error; +use crate::output::OutputInfo; +use crate::screencopy::FrameCopy; + +/// Ways to say how a region for a screenshot should be captured. +pub enum RegionCapturer { + /// Capture all of the given outputs. + Outputs(Vec), + /// Capture an already known `LogicalRegion`. + Region(LogicalRegion), +} + +/// `Region` where the coordinate system is the logical coordinate system used +/// in Wayland to position outputs. Top left is (0, 0) and any transforms and +/// scaling have been applied. +#[derive(Debug, Copy, Clone)] +pub struct LogicalRegion { + pub inner: Region, +} + +/// An embedded region is a region entirely inside of another (often an output). +/// +/// It can only be contained inside of another and cannot exceed its bounds. +/// +/// Example of what +/// +/// ┌─────────────┐ +/// │ │ +/// │ ┌──────────┼──────┐ +/// │ │ │ ├──► Viewport +/// │ │ │ │ +/// │ │ ├──────┼─────────────────┐ +/// │ │ │xxxxxx│ │ +/// │ │ │xxxxx◄├─── Inner region │ +/// │ └──────────┼──────┘ │ +/// │ │ │ +/// │ │ Screen 2 ├──► Relative to +/// │ ├────────────────────────┘ +/// │ │ +/// │ Screen 1 │ +/// └─────────────┘ +#[derive(Debug, Copy, Clone)] +pub struct EmbeddedRegion { + /// The coordinate sysd + pub relative_to: LogicalRegion, + pub inner: Region, +} + +/// Rectangle area in an unspecified coordinate system. +/// +/// Use `LogicalRegion` or `EmbeddedRegion` instead as they convey the +/// coordinate system used. +#[derive(Debug, Copy, Clone)] +pub struct Region { + /// X coordinate of the area to capture. + pub x: i32, + /// y coordinate of the area to capture. + pub y: i32, + /// Width of the capture area. + pub width: i32, + /// Height of the capture area. + pub height: i32, +} + +impl EmbeddedRegion { + /// Given two `LogicalRegion`s, one seen as the `viewport` and the other + /// `relative_to` (think the output we want to capture), create an + /// embedded region that is entirely inside of the `relative_to` region. + /// + /// See `EmbeddedRegion` for an example ASCII visualisation. + #[tracing::instrument(ret, level = "debug")] + pub fn new(viewport: LogicalRegion, relative_to: LogicalRegion) -> Option { + let x_relative: i32 = viewport.inner.x - relative_to.inner.x; + let y_relative = viewport.inner.y - relative_to.inner.y; + + let x1 = cmp::max(x_relative, 0); + let x2 = cmp::min(x_relative + viewport.inner.width, relative_to.inner.width); + let width = x2 - x1; + if width <= 0 { + return None; + } + + let y1 = cmp::max(y_relative, 0); + let y2 = cmp::min(y_relative + viewport.inner.height, relative_to.inner.height); + let height = y2 - y1; + if height <= 0 { + return None; + } + + Some(Self { + relative_to: relative_to, + inner: Region { + x: x1, + y: y1, + width, + height, + }, + }) + } + + /// Return the `LogicalRegion` of the embedded region. + /// + /// Note that this remains a region of the same size, it's not the inverse + /// of `EmbeddedRegion::new` which removes the parts that are outside of + /// the `relative_to` region. + pub fn logical(&self) -> LogicalRegion { + LogicalRegion { + inner: Region { + x: self.relative_to.inner.x as i32 + self.inner.x, + y: self.relative_to.inner.y as i32 + self.inner.y, + width: self.inner.width, + height: self.inner.height, + }, + } + } +} + +impl std::fmt::Display for EmbeddedRegion { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "{region} relative to {relative_to}", + region = self.inner, + relative_to = self.relative_to, + ) + } +} + +impl std::fmt::Display for Region { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "({x}, {y}) ({width}x{height})", + x = self.x, + y = self.y, + width = self.width, + height = self.height, + ) + } +} + +impl std::fmt::Display for LogicalRegion { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{inner}", inner = self.inner) + } +} + +impl From<&OutputInfo> for LogicalRegion { + fn from(output_info: &OutputInfo) -> Self { + LogicalRegion { + inner: Region { + x: output_info.dimensions.x, + y: output_info.dimensions.y, + width: output_info.dimensions.width, + height: output_info.dimensions.height, + }, + } + } +} + +impl TryFrom<&Vec> for LogicalRegion { + type Error = Error; + + fn try_from(output_info: &Vec) -> std::result::Result { + let x1 = output_info + .iter() + .map(|output| output.dimensions.x) + .min() + .unwrap(); + let y1 = output_info + .iter() + .map(|output| output.dimensions.y) + .min() + .unwrap(); + let x2 = output_info + .iter() + .map(|output| output.dimensions.x + output.dimensions.width) + .max() + .unwrap(); + let y2 = output_info + .iter() + .map(|output| output.dimensions.y + output.dimensions.height) + .max() + .unwrap(); + Ok(LogicalRegion { + inner: Region { + x: x1, + y: y1, + width: x2 - x1, + height: y2 - y1, + }, + }) + } +} + +impl From<&FrameCopy> for LogicalRegion { + fn from(frame_copy: &FrameCopy) -> Self { + let (width, height) = ( + frame_copy.frame_format.width as i32, + frame_copy.frame_format.height as i32, + ); + let is_portait = match frame_copy.transform { + Transform::_90 | Transform::_270 | Transform::Flipped90 | Transform::Flipped270 => true, + _ => false, + }; + LogicalRegion { + inner: Region { + x: frame_copy.position.0 as i32, + y: frame_copy.position.1 as i32, + width: if is_portait { height } else { width }, + height: if is_portait { width } else { height }, + }, + } + } +} diff --git a/libwayshot/src/screencopy.rs b/libwayshot/src/screencopy.rs index e2713582..2df15601 100644 --- a/libwayshot/src/screencopy.rs +++ b/libwayshot/src/screencopy.rs @@ -11,20 +11,38 @@ use nix::{ sys::{memfd, mman, stat}, unistd, }; -use wayland_client::protocol::{wl_output, wl_shm::Format}; +use wayland_client::protocol::{ + wl_buffer::WlBuffer, wl_output, wl_shm::Format, wl_shm_pool::WlShmPool, +}; use crate::{Error, Result}; +pub struct FrameGuard { + pub buffer: WlBuffer, + pub shm_pool: WlShmPool, +} + +impl Drop for FrameGuard { + fn drop(&mut self) { + self.buffer.destroy(); + self.shm_pool.destroy(); + } +} + /// Type of frame supported by the compositor. For now we only support Argb8888, Xrgb8888, and /// Xbgr8888. +/// +/// See `zwlr_screencopy_frame_v1::Event::Buffer` as it's retrieved from there. #[derive(Debug, Copy, Clone, PartialEq)] pub struct FrameFormat { pub format: Format, pub width: u32, pub height: u32, + /// Stride is the number of bytes between the start of a row and the start of the next row. pub stride: u32, } +#[tracing::instrument(skip(frame_mmap))] fn create_image_buffer

( frame_format: &FrameFormat, frame_mmap: &MmapMut, @@ -32,6 +50,7 @@ fn create_image_buffer

( where P: Pixel, { + tracing::debug!("Creating image buffer"); ImageBuffer::from_vec(frame_format.width, frame_format.height, frame_mmap.to_vec()) .ok_or(Error::BufferTooSmall) } @@ -44,12 +63,13 @@ pub struct FrameCopy { pub frame_color_type: ColorType, pub frame_mmap: MmapMut, pub transform: wl_output::Transform, + pub position: (i64, i64), } -impl TryFrom for DynamicImage { +impl TryFrom<&FrameCopy> for DynamicImage { type Error = Error; - fn try_from(value: FrameCopy) -> Result { + fn try_from(value: &FrameCopy) -> Result { Ok(match value.frame_color_type { ColorType::Rgb8 => { Self::ImageRgb8(create_image_buffer(&value.frame_format, &value.frame_mmap)?) diff --git a/wayshot/src/utils.rs b/wayshot/src/utils.rs index e24caf74..c6cd60fa 100644 --- a/wayshot/src/utils.rs +++ b/wayshot/src/utils.rs @@ -3,9 +3,9 @@ use std::{ time::{SystemTime, UNIX_EPOCH}, }; -use libwayshot::CaptureRegion; +use libwayshot::region::{LogicalRegion, Region}; -pub fn parse_geometry(g: &str) -> Option { +pub fn parse_geometry(g: &str) -> Option { let tail = g.trim(); let x_coordinate: i32; let y_coordinate: i32; @@ -32,11 +32,13 @@ pub fn parse_geometry(g: &str) -> Option { height = tail.parse::().ok()?; } - Some(CaptureRegion { - x_coordinate, - y_coordinate, - width, - height, + Some(LogicalRegion { + inner: Region { + x: x_coordinate, + y: y_coordinate, + width, + height, + }, }) } From a023ff1ce2daeb389006489afb3b320edf3637b8 Mon Sep 17 00:00:00 2001 From: Andreas Backx Date: Fri, 2 Feb 2024 18:59:36 +0000 Subject: [PATCH 2/6] Add freeze functionality --- Cargo.lock | 17 ++++++ libwayshot/src/dispatch.rs | 69 +++++++++++++++++----- libwayshot/src/error.rs | 2 + libwayshot/src/lib.rs | 113 ++++++++++++++++++++++++++++++++++++- libwayshot/src/output.rs | 7 ++- libwayshot/src/region.rs | 6 +- wayshot/Cargo.toml | 1 + wayshot/src/clap.rs | 4 +- wayshot/src/utils.rs | 36 ++++++------ wayshot/src/wayshot.rs | 29 ++++++---- 10 files changed, 236 insertions(+), 48 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 2d5bff5d..500045a8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -207,6 +207,16 @@ dependencies = [ "windows-sys 0.52.0", ] +[[package]] +name = "eyre" +version = "0.6.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4c2b6b5a29c02cdc822728b7d7b8ae1bab3e3b05d44522770ddd49722eeac7eb" +dependencies = [ + "indenter", + "once_cell", +] + [[package]] name = "fastrand" version = "2.0.1" @@ -257,6 +267,12 @@ dependencies = [ "qoi", ] +[[package]] +name = "indenter" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ce23b50ad8242c51a442f3ff322d56b02f08852c77e4c0b4d3fd684abc89c683" + [[package]] name = "jpeg-decoder" version = "0.3.0" @@ -756,6 +772,7 @@ version = "1.3.2-dev" dependencies = [ "clap", "dialoguer", + "eyre", "flate2", "image", "libwayshot", diff --git a/libwayshot/src/dispatch.rs b/libwayshot/src/dispatch.rs index 170e1026..f4eedb06 100644 --- a/libwayshot/src/dispatch.rs +++ b/libwayshot/src/dispatch.rs @@ -1,4 +1,5 @@ use std::{ + collections::HashSet, process::exit, sync::atomic::{AtomicBool, Ordering}, }; @@ -6,8 +7,9 @@ use wayland_client::{ delegate_noop, globals::GlobalListContents, protocol::{ - wl_buffer::WlBuffer, wl_output, wl_output::WlOutput, wl_registry, wl_registry::WlRegistry, - wl_shm::WlShm, wl_shm_pool::WlShmPool, + wl_buffer::WlBuffer, wl_compositor::WlCompositor, wl_output, wl_output::WlOutput, + wl_registry, wl_registry::WlRegistry, wl_shm::WlShm, wl_shm_pool::WlShmPool, + wl_surface::WlSurface, }, Connection, Dispatch, QueueHandle, WEnum, WEnum::Value, @@ -15,6 +17,10 @@ use wayland_client::{ use wayland_protocols::xdg::xdg_output::zv1::client::{ zxdg_output_manager_v1::ZxdgOutputManagerV1, zxdg_output_v1, zxdg_output_v1::ZxdgOutputV1, }; +use wayland_protocols_wlr::layer_shell::v1::client::{ + zwlr_layer_shell_v1::ZwlrLayerShellV1, + zwlr_layer_surface_v1::{self, ZwlrLayerSurfaceV1}, +}; use wayland_protocols_wlr::screencopy::v1::client::{ zwlr_screencopy_frame_v1, zwlr_screencopy_frame_v1::ZwlrScreencopyFrameV1, zwlr_screencopy_manager_v1::ZwlrScreencopyManagerV1, @@ -58,16 +64,9 @@ impl Dispatch for OutputCaptureState { name: "".to_string(), description: String::new(), transform: wl_output::Transform::Normal, - dimensions: OutputPositioning { - x: 0, - y: 0, - width: 0, - height: 0, - }, - mode: WlOutputMode { - width: 0, - height: 0, - }, + scale: 1, + dimensions: OutputPositioning::default(), + mode: WlOutputMode::default(), }); } else { tracing::error!("Ignoring a wl_output with version < 4."); @@ -109,7 +108,11 @@ impl Dispatch for OutputCaptureState { } => { output.transform = transform; } - _ => (), + wl_output::Event::Scale { factor } => { + output.scale = factor; + } + wl_output::Event::Done => {} + _ => {} } } } @@ -227,3 +230,43 @@ impl wayland_client::Dispatch for W ) { } } + +pub struct LayerShellState { + pub configured_outputs: HashSet, +} + +delegate_noop!(LayerShellState: ignore WlCompositor); +delegate_noop!(LayerShellState: ignore WlShm); +delegate_noop!(LayerShellState: ignore WlShmPool); +delegate_noop!(LayerShellState: ignore WlBuffer); +delegate_noop!(LayerShellState: ignore ZwlrLayerShellV1); +delegate_noop!(LayerShellState: ignore WlSurface); + +impl wayland_client::Dispatch for LayerShellState { + // No need to instrument here, span from lib.rs is automatically used. + fn event( + state: &mut Self, + proxy: &ZwlrLayerSurfaceV1, + event: ::Event, + data: &WlOutput, + _conn: &Connection, + _qhandle: &QueueHandle, + ) { + match event { + zwlr_layer_surface_v1::Event::Configure { + serial, + width: _, + height: _, + } => { + tracing::debug!("Acking configure"); + state.configured_outputs.insert(data.clone()); + proxy.ack_configure(serial); + tracing::trace!("Acked configure"); + } + zwlr_layer_surface_v1::Event::Closed => { + tracing::debug!("Closed") + } + _ => {} + } + } +} diff --git a/libwayshot/src/error.rs b/libwayshot/src/error.rs index ddeff581..ab0e1ec3 100644 --- a/libwayshot/src/error.rs +++ b/libwayshot/src/error.rs @@ -27,4 +27,6 @@ pub enum Error { NoSupportedBufferFormat, #[error("Cannot find required wayland protocol")] ProtocolNotFound(String), + #[error("error occurred in freeze callback")] + FreezeCallbackError, } diff --git a/libwayshot/src/lib.rs b/libwayshot/src/lib.rs index c93cb335..2128ed55 100644 --- a/libwayshot/src/lib.rs +++ b/libwayshot/src/lib.rs @@ -12,6 +12,7 @@ pub mod region; mod screencopy; use std::{ + collections::HashSet, fs::File, os::fd::AsFd, process::exit, @@ -19,6 +20,7 @@ use std::{ thread, }; +use dispatch::LayerShellState; use image::{imageops::replace, DynamicImage}; use memmap2::MmapMut; use region::{EmbeddedRegion, RegionCapturer}; @@ -27,6 +29,7 @@ use tracing::debug; use wayland_client::{ globals::{registry_queue_init, GlobalList}, protocol::{ + wl_compositor::WlCompositor, wl_output::WlOutput, wl_shm::{self, WlShm}, }, @@ -35,9 +38,15 @@ use wayland_client::{ use wayland_protocols::xdg::xdg_output::zv1::client::{ zxdg_output_manager_v1::ZxdgOutputManagerV1, zxdg_output_v1::ZxdgOutputV1, }; -use wayland_protocols_wlr::screencopy::v1::client::{ - zwlr_screencopy_frame_v1::ZwlrScreencopyFrameV1, - zwlr_screencopy_manager_v1::ZwlrScreencopyManagerV1, +use wayland_protocols_wlr::{ + layer_shell::v1::client::{ + zwlr_layer_shell_v1::{Layer, ZwlrLayerShellV1}, + zwlr_layer_surface_v1::Anchor, + }, + screencopy::v1::client::{ + zwlr_screencopy_frame_v1::ZwlrScreencopyFrameV1, + zwlr_screencopy_manager_v1::ZwlrScreencopyManagerV1, + }, }; use crate::{ @@ -406,6 +415,87 @@ impl WayshotConnection { Ok(frame_copies) } + fn overlay_frames(&self, frames: &Vec<(FrameCopy, FrameGuard, OutputInfo)>) -> Result<()> { + let mut state = LayerShellState { + configured_outputs: HashSet::new(), + }; + let mut event_queue: EventQueue = + self.conn.new_event_queue::(); + let qh = event_queue.handle(); + + let compositor = match self.globals.bind::(&qh, 3..=3, ()) { + Ok(x) => x, + Err(e) => { + tracing::error!( + "Failed to create compositor Does your compositor implement WlCompositor?" + ); + tracing::error!("err: {e}"); + return Err(Error::ProtocolNotFound( + "WlCompositor not found".to_string(), + )); + } + }; + let layer_shell = match self.globals.bind::(&qh, 1..=1, ()) { + Ok(x) => x, + Err(e) => { + tracing::error!( + "Failed to create layer shell. Does your compositor implement WlrLayerShellV1?" + ); + tracing::error!("err: {e}"); + return Err(Error::ProtocolNotFound( + "WlrLayerShellV1 not found".to_string(), + )); + } + }; + + for (frame_copy, frame_guard, output_info) in frames { + tracing::span!( + tracing::Level::DEBUG, + "overlay_frames::surface", + output = output_info.name.as_str() + ) + .in_scope(|| -> Result<()> { + let surface = compositor.create_surface(&qh, ()); + + let layer_surface = layer_shell.get_layer_surface( + &surface, + Some(&output_info.wl_output), + Layer::Top, + "wayshot".to_string(), + &qh, + output_info.wl_output.clone(), + ); + + layer_surface.set_exclusive_zone(-1); + layer_surface.set_anchor(Anchor::Top | Anchor::Left); + layer_surface.set_size( + frame_copy.frame_format.width, + frame_copy.frame_format.height, + ); + + debug!("Committing surface creation changes."); + surface.commit(); + + debug!("Waiting for layer surface to be configured."); + while !state.configured_outputs.contains(&output_info.wl_output) { + event_queue.blocking_dispatch(&mut state)?; + } + + surface.set_buffer_transform(output_info.transform); + surface.set_buffer_scale(output_info.scale); + surface.attach(Some(&frame_guard.buffer), 0, 0); + + debug!("Committing surface with attached buffer."); + surface.commit(); + + event_queue.blocking_dispatch(&mut state)?; + + Ok(()) + })?; + } + Ok(()) + } + /// Take a screenshot from the specified region. fn screenshot_region_capturer( &self, @@ -445,6 +535,11 @@ impl WayshotConnection { }) }) .collect(), + RegionCapturer::Freeze(_) => self + .get_all_outputs() + .into_iter() + .map(|output_info| (output_info.clone(), None)) + .collect(), }; let frames = self.capture_frame_copies(outputs_capture_regions, cursor_overlay)?; @@ -452,6 +547,9 @@ impl WayshotConnection { let capture_region: LogicalRegion = match region_capturer { RegionCapturer::Outputs(ref outputs) => outputs.try_into()?, RegionCapturer::Region(region) => region, + RegionCapturer::Freeze(callback) => { + self.overlay_frames(&frames).and_then(|_| callback())? + } }; thread::scope(|scope| { @@ -529,6 +627,15 @@ impl WayshotConnection { self.screenshot_region_capturer(RegionCapturer::Region(capture_region), cursor_overlay) } + /// Take a screenshot, overlay the screenshot, run the callback, and then + /// unfreeze the screenshot and return the selected region. + pub fn screenshot_freeze( + &self, + callback: Box Result>, + cursor_overlay: bool, + ) -> Result { + self.screenshot_region_capturer(RegionCapturer::Freeze(callback), cursor_overlay) + } /// shot one ouput pub fn screenshot_single_output( &self, diff --git a/libwayshot/src/output.rs b/libwayshot/src/output.rs index ccca1c80..986f74df 100644 --- a/libwayshot/src/output.rs +++ b/libwayshot/src/output.rs @@ -3,23 +3,24 @@ use wayland_client::protocol::{wl_output, wl_output::WlOutput}; /// Represents an accessible wayland output. /// /// Do not instantiate, instead use [`crate::WayshotConnection::get_all_outputs`]. -#[derive(Debug, Clone)] +#[derive(Debug, Clone, PartialEq, Eq, Hash)] pub struct OutputInfo { pub wl_output: WlOutput, pub name: String, pub description: String, pub transform: wl_output::Transform, + pub scale: i32, pub dimensions: OutputPositioning, pub mode: WlOutputMode, } -#[derive(Default, Debug, Clone, PartialEq, Eq)] +#[derive(Default, Debug, Clone, PartialEq, Eq, Hash)] pub struct WlOutputMode { pub width: i32, pub height: i32, } -#[derive(Default, Debug, Clone, PartialEq, Eq)] +#[derive(Default, Debug, Clone, PartialEq, Eq, Hash)] pub struct OutputPositioning { pub x: i32, pub y: i32, diff --git a/libwayshot/src/region.rs b/libwayshot/src/region.rs index a82e0cdc..490e6082 100644 --- a/libwayshot/src/region.rs +++ b/libwayshot/src/region.rs @@ -2,7 +2,7 @@ use std::cmp; use wayland_client::protocol::wl_output::Transform; -use crate::error::Error; +use crate::error::{Error, Result}; use crate::output::OutputInfo; use crate::screencopy::FrameCopy; @@ -12,6 +12,10 @@ pub enum RegionCapturer { Outputs(Vec), /// Capture an already known `LogicalRegion`. Region(LogicalRegion), + /// The outputs will be "frozen" to the user at which point the given + /// callback is called to get the region to capture. This callback is often + /// a user interaction to let the user select a region. + Freeze(Box Result>), } /// `Region` where the coordinate system is the logical coordinate system used diff --git a/wayshot/Cargo.toml b/wayshot/Cargo.toml index 8d4aa528..e4ac7af3 100644 --- a/wayshot/Cargo.toml +++ b/wayshot/Cargo.toml @@ -29,6 +29,7 @@ image = { version = "0.24", default-features = false, features = [ ] } dialoguer = { version = "0.11.0", features = ["fuzzy-select"] } +eyre = "0.6.8" [[bin]] name = "wayshot" diff --git a/wayshot/src/clap.rs b/wayshot/src/clap.rs index 37319334..dfed67b4 100644 --- a/wayshot/src/clap.rs +++ b/wayshot/src/clap.rs @@ -12,10 +12,10 @@ pub fn set_flags() -> Command { .help("Enable debug mode"), ) .arg( - arg!(-s --slurp ) + arg!(-s --slurp ) .required(false) .action(ArgAction::Set) - .help("Choose a portion of your display to screenshot using slurp"), + .help("Arguments to call slurp with for selecting a region"), ) .arg( arg!(-f - -file ) diff --git a/wayshot/src/utils.rs b/wayshot/src/utils.rs index c6cd60fa..cea4217d 100644 --- a/wayshot/src/utils.rs +++ b/wayshot/src/utils.rs @@ -1,3 +1,4 @@ +use eyre::{ContextCompat, Result}; use std::{ process::exit, time::{SystemTime, UNIX_EPOCH}, @@ -5,34 +6,37 @@ use std::{ use libwayshot::region::{LogicalRegion, Region}; -pub fn parse_geometry(g: &str) -> Option { +pub fn parse_geometry(g: &str) -> Result { let tail = g.trim(); let x_coordinate: i32; let y_coordinate: i32; let width: i32; let height: i32; + let validation_error = + "Invalid geometry provided.\nValid geometries:\n1) %d,%d %dx%d\n2) %d %d %d %d"; + if tail.contains(',') { // this accepts: "%d,%d %dx%d" - let (head, tail) = tail.split_once(',')?; - x_coordinate = head.parse::().ok()?; - let (head, tail) = tail.split_once(' ')?; - y_coordinate = head.parse::().ok()?; - let (head, tail) = tail.split_once('x')?; - width = head.parse::().ok()?; - height = tail.parse::().ok()?; + let (head, tail) = tail.split_once(',').wrap_err(validation_error)?; + x_coordinate = head.parse::()?; + let (head, tail) = tail.split_once(' ').wrap_err(validation_error)?; + y_coordinate = head.parse::()?; + let (head, tail) = tail.split_once('x').wrap_err(validation_error)?; + width = head.parse::()?; + height = tail.parse::()?; } else { // this accepts: "%d %d %d %d" - let (head, tail) = tail.split_once(' ')?; - x_coordinate = head.parse::().ok()?; - let (head, tail) = tail.split_once(' ')?; - y_coordinate = head.parse::().ok()?; - let (head, tail) = tail.split_once(' ')?; - width = head.parse::().ok()?; - height = tail.parse::().ok()?; + let (head, tail) = tail.split_once(' ').wrap_err(validation_error)?; + x_coordinate = head.parse::()?; + let (head, tail) = tail.split_once(' ').wrap_err(validation_error)?; + y_coordinate = head.parse::()?; + let (head, tail) = tail.split_once(' ').wrap_err(validation_error)?; + width = head.parse::()?; + height = tail.parse::()?; } - Some(LogicalRegion { + Ok(LogicalRegion { inner: Region { x: x_coordinate, y: y_coordinate, diff --git a/wayshot/src/wayshot.rs b/wayshot/src/wayshot.rs index b4b3c01f..4353d579 100644 --- a/wayshot/src/wayshot.rs +++ b/wayshot/src/wayshot.rs @@ -1,10 +1,10 @@ use std::{ - error::Error, io::{stdout, BufWriter, Cursor, Write}, - process::exit, + process::{exit, Command}, }; -use libwayshot::WayshotConnection; +use eyre::Result; +use libwayshot::{region::LogicalRegion, WayshotConnection}; mod clap; mod utils; @@ -29,7 +29,7 @@ where Some(selection) } -fn main() -> Result<(), Box> { +fn main() -> Result<()> { let args = clap::set_flags().get_matches(); let level = if args.get_flag("debug") { Level::TRACE @@ -86,12 +86,21 @@ fn main() -> Result<(), Box> { } let image_buffer = if let Some(slurp_region) = args.get_one::("slurp") { - if let Some(region) = utils::parse_geometry(slurp_region) { - wayshot_conn.screenshot(region, cursor_overlay)? - } else { - tracing::error!("Invalid geometry specification"); - exit(1); - } + let slurp_region = slurp_region.clone(); + wayshot_conn.screenshot_freeze( + Box::new(move || { + || -> Result { + let slurp_output = Command::new("slurp") + .args(slurp_region.split(" ")) + .output()? + .stdout; + + utils::parse_geometry(&String::from_utf8(slurp_output)?) + }() + .map_err(|_| libwayshot::Error::FreezeCallbackError) + }), + cursor_overlay, + )? } else if let Some(output_name) = args.get_one::("output") { let outputs = wayshot_conn.get_all_outputs(); if let Some(output) = outputs.iter().find(|output| &output.name == output_name) { From 87be331d815ef3176ae131b45d3353604fd87601 Mon Sep 17 00:00:00 2001 From: Andreas Backx Date: Fri, 2 Feb 2024 19:01:37 +0000 Subject: [PATCH 3/6] Move to clap_derive so CLI arguments are typed * Renames `clap.rs` to `cli.rs` because otherwise there's confusion between the `clap` crate and local module. * `Cli` struct added that almost identically represents the current state of the CLI with no logical changes. --- ## `--help` Comparison Before: https://gist.github.com/AndreasBackx/5945b366e989159f4669e7ba30c13239 After: https://gist.github.com/AndreasBackx/8929c8bde080eac0cafd33128210b0cc Diff (ignoring whitespace changes due to table alignment): ```diff 1c1 < Screenshot tool for compositors implementing zwlr_screencopy_v1. --- > Screenshot tool for wlroots based compositors implementing the zwlr_screencopy_v1 protocol. 9d8 < -c, --cursor Enable cursor in screenshots 10a10 > -c, --cursor Enable cursor in screenshots 12c12 < -l, --listoutputs List all valid outputs --- > -l, --list-outputs List all valid outputs 14c14 < --chooseoutput Present a fuzzy selector for outputs --- > --choose-output Present a fuzzy selector for outputs 16a17 > ``` You can see that the only changes are: - About is longer, this is now using the value from Cargo.toml instead of a duplicate text that was shorter. - Some have a dash where the English words would have a space, e.g: "list-outputs" instead of "listoutputs". I've also made the old still work via an alias: https://gist.github.com/AndreasBackx/6025e91844e3d766d4264a01ae4d1a71 This seems like a tiny improvement? I plan to make further changes later, but I want to keep PRs separate. --- Cargo.lock | 31 +++++++++++++++---- wayshot/Cargo.toml | 2 +- wayshot/src/clap.rs | 67 ------------------------------------------ wayshot/src/cli.rs | 43 +++++++++++++++++++++++++++ wayshot/src/wayshot.rs | 40 ++++++++++--------------- 5 files changed, 85 insertions(+), 98 deletions(-) delete mode 100644 wayshot/src/clap.rs create mode 100644 wayshot/src/cli.rs diff --git a/Cargo.lock b/Cargo.lock index 500045a8..9244fa4d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10,9 +10,9 @@ checksum = "f26201604c87b1e01bd3d98f8d5d9a8fcbb815e8cedb41ffccbeb4bf593a35fe" [[package]] name = "anstream" -version = "0.6.5" +version = "0.6.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d664a92ecae85fd0a7392615844904654d1d5f5514837f471ddef4a057aba1b6" +checksum = "6e2e1ebcb11de5c03c67de28a7df593d32191b44939c482e97702baaaa6ab6a5" dependencies = [ "anstyle", "anstyle-parse", @@ -103,18 +103,19 @@ checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd" [[package]] name = "clap" -version = "4.4.11" +version = "4.4.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bfaff671f6b22ca62406885ece523383b9b64022e341e53e009a62ebc47a45f2" +checksum = "1e578d6ec4194633722ccf9544794b71b1385c3c027efe0c55db226fc880865c" dependencies = [ "clap_builder", + "clap_derive", ] [[package]] name = "clap_builder" -version = "4.4.11" +version = "4.4.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a216b506622bb1d316cd51328dce24e07bdff4a6128a47c7e7fad11878d5adbb" +checksum = "4df4df40ec50c46000231c914968278b1eb05098cf8f1b3a518a95030e71d1c7" dependencies = [ "anstream", "anstyle", @@ -122,6 +123,18 @@ dependencies = [ "strsim", ] +[[package]] +name = "clap_derive" +version = "4.4.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf9804afaaf59a91e75b022a30fb7229a7901f60c755489cc61c9b423b836442" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "clap_lex" version = "0.6.0" @@ -251,6 +264,12 @@ dependencies = [ "thread_local", ] +[[package]] +name = "heck" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "95505c38b4572b2d910cecb0281560f54b440a19336cbbcb27bf6ce6adc6f5a8" + [[package]] name = "image" version = "0.24.7" diff --git a/wayshot/Cargo.toml b/wayshot/Cargo.toml index e4ac7af3..344654f9 100644 --- a/wayshot/Cargo.toml +++ b/wayshot/Cargo.toml @@ -18,7 +18,7 @@ tracing.workspace = true libwayshot.workspace = true -clap = "4.4.6" +clap = { version = "4.4.18", features = ["derive"] } tracing-subscriber = "0.3.17" image = { version = "0.24", default-features = false, features = [ diff --git a/wayshot/src/clap.rs b/wayshot/src/clap.rs deleted file mode 100644 index dfed67b4..00000000 --- a/wayshot/src/clap.rs +++ /dev/null @@ -1,67 +0,0 @@ -use clap::{arg, ArgAction, Command}; - -pub fn set_flags() -> Command { - Command::new("wayshot") - .version(env!("CARGO_PKG_VERSION")) - .author(env!("CARGO_PKG_AUTHORS")) - .about("Screenshot tool for compositors implementing zwlr_screencopy_v1.") - .arg( - arg!(-d - -debug) - .required(false) - .action(ArgAction::SetTrue) - .help("Enable debug mode"), - ) - .arg( - arg!(-s --slurp ) - .required(false) - .action(ArgAction::Set) - .help("Arguments to call slurp with for selecting a region"), - ) - .arg( - arg!(-f - -file ) - .required(false) - .conflicts_with("stdout") - .action(ArgAction::Set) - .help("Mention a custom file path"), - ) - .arg( - arg!(-c - -cursor) - .required(false) - .action(ArgAction::SetTrue) - .help("Enable cursor in screenshots"), - ) - .arg( - arg!(--stdout) - .required(false) - .conflicts_with("file") - .action(ArgAction::SetTrue) - .help("Output the image data to standard out"), - ) - .arg( - arg!(-e --extension ) - .required(false) - .action(ArgAction::Set) - .help("Set image encoder (Png is default)"), - ) - .arg( - arg!(-l - -listoutputs) - .required(false) - .action(ArgAction::SetTrue) - .help("List all valid outputs"), - ) - .arg( - arg!(-o --output ) - .required(false) - .action(ArgAction::Set) - .conflicts_with("slurp") - .help("Choose a particular display to screenshot"), - ) - .arg( - arg!(--chooseoutput) - .required(false) - .action(ArgAction::SetTrue) - .conflicts_with("slurp") - .conflicts_with("output") - .help("Present a fuzzy selector for outputs"), - ) -} diff --git a/wayshot/src/cli.rs b/wayshot/src/cli.rs new file mode 100644 index 00000000..182737b2 --- /dev/null +++ b/wayshot/src/cli.rs @@ -0,0 +1,43 @@ +use clap::arg; + +use clap::Parser; + +#[derive(Parser)] +#[command(version, about)] +pub struct Cli { + /// Enable debug mode + #[arg(short, long, action = clap::ArgAction::SetTrue)] + pub debug: bool, + + /// Arguments to call slurp with for selecting a region + #[arg(short, long, value_name = "SLURP_ARGS")] + pub slurp: Option, + + /// Mention a custom file path + #[arg(short, long, conflicts_with = "stdout", value_name = "FILE_PATH")] + pub file: Option, + + /// Output the image data to standard out + #[arg(long, conflicts_with = "file", action = clap::ArgAction::SetTrue)] + pub stdout: bool, + + /// Enable cursor in screenshots + #[arg(short, long, action = clap::ArgAction::SetTrue)] + pub cursor: bool, + + /// Set image encoder (Png is default) + #[arg(short, long, value_name = "FILE_EXTENSION")] + pub extension: Option, + + /// List all valid outputs + #[arg(short, long, alias = "listoutputs", action = clap::ArgAction::SetTrue)] + pub list_outputs: bool, + + /// Choose a particular display to screenshot + #[arg(short, long, conflicts_with = "slurp")] + pub output: Option, + + /// Present a fuzzy selector for outputs + #[arg(long, alias = "chooseoutput", conflicts_with_all = ["slurp", "output"], action = clap::ArgAction::SetTrue)] + pub choose_output: bool, +} diff --git a/wayshot/src/wayshot.rs b/wayshot/src/wayshot.rs index 4353d579..8e3bc8af 100644 --- a/wayshot/src/wayshot.rs +++ b/wayshot/src/wayshot.rs @@ -3,10 +3,11 @@ use std::{ process::{exit, Command}, }; +use clap::Parser; use eyre::Result; use libwayshot::{region::LogicalRegion, WayshotConnection}; -mod clap; +mod cli; mod utils; use dialoguer::{theme::ColorfulTheme, FuzzySelect}; @@ -30,18 +31,14 @@ where } fn main() -> Result<()> { - let args = clap::set_flags().get_matches(); - let level = if args.get_flag("debug") { - Level::TRACE - } else { - Level::INFO - }; + let cli = cli::Cli::parse(); + let level = if cli.debug { Level::TRACE } else { Level::INFO }; tracing_subscriber::fmt() .with_max_level(level) .with_writer(std::io::stderr) .init(); - let extension = if let Some(extension) = args.get_one::("extension") { + let extension = if let Some(extension) = cli.extension { let ext = extension.trim().to_lowercase(); tracing::debug!("Using custom extension: {:#?}", ext); @@ -62,9 +59,9 @@ fn main() -> Result<()> { let mut file_is_stdout: bool = false; let mut file_path: Option = None; - if args.get_flag("stdout") { + if cli.stdout { file_is_stdout = true; - } else if let Some(filepath) = args.get_one::("file") { + } else if let Some(filepath) = cli.file { file_path = Some(filepath.trim().to_string()); } else { file_path = Some(utils::get_default_file_name(extension)); @@ -72,7 +69,7 @@ fn main() -> Result<()> { let wayshot_conn = WayshotConnection::new()?; - if args.get_flag("listoutputs") { + if cli.list_outputs { let valid_outputs = wayshot_conn.get_all_outputs(); for output in valid_outputs { tracing::info!("{:#?}", output.name); @@ -80,12 +77,7 @@ fn main() -> Result<()> { exit(1); } - let mut cursor_overlay = false; - if args.get_flag("cursor") { - cursor_overlay = true; - } - - let image_buffer = if let Some(slurp_region) = args.get_one::("slurp") { + let image_buffer = if let Some(slurp_region) = cli.slurp { let slurp_region = slurp_region.clone(); wayshot_conn.screenshot_freeze( Box::new(move || { @@ -99,30 +91,30 @@ fn main() -> Result<()> { }() .map_err(|_| libwayshot::Error::FreezeCallbackError) }), - cursor_overlay, + cli.cursor, )? - } else if let Some(output_name) = args.get_one::("output") { + } else if let Some(output_name) = cli.output { let outputs = wayshot_conn.get_all_outputs(); - if let Some(output) = outputs.iter().find(|output| &output.name == output_name) { - wayshot_conn.screenshot_single_output(output, cursor_overlay)? + if let Some(output) = outputs.iter().find(|output| output.name == output_name) { + wayshot_conn.screenshot_single_output(output, cli.cursor)? } else { tracing::error!("No output found!\n"); exit(1); } - } else if args.get_flag("chooseoutput") { + } else if cli.choose_output { let outputs = wayshot_conn.get_all_outputs(); let output_names: Vec = outputs .iter() .map(|display| display.name.to_string()) .collect(); if let Some(index) = select_ouput(&output_names) { - wayshot_conn.screenshot_single_output(&outputs[index], cursor_overlay)? + wayshot_conn.screenshot_single_output(&outputs[index], cli.cursor)? } else { tracing::error!("No output found!\n"); exit(1); } } else { - wayshot_conn.screenshot_all(cursor_overlay)? + wayshot_conn.screenshot_all(cli.cursor)? }; if file_is_stdout { From 736b6515c86f5443e09aa888629c439cff1d5618 Mon Sep 17 00:00:00 2001 From: Andreas Backx Date: Fri, 2 Feb 2024 19:01:37 +0000 Subject: [PATCH 4/6] Improve CLI design This "improves" (and that is subjective) the design of the CLI. I am aiming to get some feedback on what people think of the new design: ``` Screenshot tool for wlroots based compositors implementing the zwlr_screencopy_v1 protocol. Usage: wayshot [OPTIONS] [OUTPUT] Arguments: [OUTPUT] Where to save the screenshot, "-" for stdout. Defaults to "$UNIX_TIMESTAMP-wayshot.$EXTENSION" Options: --log-level Log level to be used for printing to stderr [default: info] [possible values: trace, debug, info, warn, error] -s, --slurp Arguments to call slurp with for selecting a region -c, --cursor Enable cursor in screenshots --encoding Set image encoder, if output file contains an extension, that will be used instead [default: png] [aliases: extension, format, output-format] Possible values: - jpg: JPG/JPEG encoder - png: PNG encoder - ppm: PPM encoder - qoi: Qut encoder -l, --list-outputs List all valid outputs -o, --output Choose a particular output/display to screenshot --choose-output Present a fuzzy selector for output/display selection -h, --help Print help (see a summary with '-h') -V, --version Print version ``` The main changes are: 1. `--debug` is now `--log-level` because this makes it easy to select more specifically what log level to use. I considered using `-v`, `-vv`... to increase verbosity but the `clap-verbosity-crate` uses `log` and not `tracing`. We could use it somewhat, but we'd pull in `log` (which seems fine) and we'd need to map from `log`'s Level to `tracing`'s Level enums (they use inverse ordering). 2. `--stdout` and `--file` has been made an optional positional argument. This because it's what other CLIs often do and I wasn't sure what to call the option otherwise because `--output` and `-O`/`-o` is often what others use but we use it here to refer to displays/monitors/Wayland outputs. This avoids that confusion hopefully and I've also clarified this in the documentation. * Additionally if a path is given, its extension will always be used. So you cannot save `jpg` to `foo.png`. Perhaps this behaviour can be changed, though I don't see a reason to support this weird edge case? When is someone saving `png` to `jpg`? 3. `--extension` is `--encoding` with aliases like `extension`. Again, let me know what you think. --- .gitignore | 5 ++++ wayshot/src/cli.rs | 40 +++++++++++++------------ wayshot/src/utils.rs | 68 +++++++++++++++++++++++++++++++++++++----- wayshot/src/wayshot.rs | 64 +++++++++++++++++++-------------------- 4 files changed, 116 insertions(+), 61 deletions(-) diff --git a/.gitignore b/.gitignore index 5a4f6354..2622eacf 100644 --- a/.gitignore +++ b/.gitignore @@ -2,3 +2,8 @@ target *.gz *.out .direnv +*.jpg +*.jpeg +*.png +*.ppm +*.qoi diff --git a/wayshot/src/cli.rs b/wayshot/src/cli.rs index 182737b2..a076af5c 100644 --- a/wayshot/src/cli.rs +++ b/wayshot/src/cli.rs @@ -1,43 +1,45 @@ +use std::path::PathBuf; + use clap::arg; use clap::Parser; +use crate::utils::EncodingFormat; +use clap::builder::TypedValueParser; + #[derive(Parser)] #[command(version, about)] pub struct Cli { - /// Enable debug mode - #[arg(short, long, action = clap::ArgAction::SetTrue)] - pub debug: bool, + /// Where to save the screenshot, "-" for stdout. Defaults to "$UNIX_TIMESTAMP-wayshot.$EXTENSION". + #[arg(value_name = "OUTPUT")] + pub file: Option, + + /// Log level to be used for printing to stderr + #[arg(long, default_value = "info", value_parser = clap::builder::PossibleValuesParser::new(["trace", "debug", "info", "warn", "error"]).map(|s| -> tracing::Level{ s.parse().unwrap()}))] + pub log_level: tracing::Level, /// Arguments to call slurp with for selecting a region #[arg(short, long, value_name = "SLURP_ARGS")] pub slurp: Option, - /// Mention a custom file path - #[arg(short, long, conflicts_with = "stdout", value_name = "FILE_PATH")] - pub file: Option, - - /// Output the image data to standard out - #[arg(long, conflicts_with = "file", action = clap::ArgAction::SetTrue)] - pub stdout: bool, - /// Enable cursor in screenshots - #[arg(short, long, action = clap::ArgAction::SetTrue)] + #[arg(short, long)] pub cursor: bool, - /// Set image encoder (Png is default) - #[arg(short, long, value_name = "FILE_EXTENSION")] - pub extension: Option, + /// Set image encoder, by default uses the file extension from the OUTPUT + /// positional argument. Otherwise defaults to png. + #[arg(long, visible_aliases = ["extension", "format", "output-format"], value_name = "FILE_EXTENSION")] + pub encoding: Option, /// List all valid outputs - #[arg(short, long, alias = "listoutputs", action = clap::ArgAction::SetTrue)] + #[arg(short, long, alias = "listoutputs")] pub list_outputs: bool, - /// Choose a particular display to screenshot + /// Choose a particular output/display to screenshot #[arg(short, long, conflicts_with = "slurp")] pub output: Option, - /// Present a fuzzy selector for outputs - #[arg(long, alias = "chooseoutput", conflicts_with_all = ["slurp", "output"], action = clap::ArgAction::SetTrue)] + /// Present a fuzzy selector for output/display selection + #[arg(long, alias = "chooseoutput", conflicts_with_all = ["slurp", "output"])] pub choose_output: bool, } diff --git a/wayshot/src/utils.rs b/wayshot/src/utils.rs index cea4217d..1e394452 100644 --- a/wayshot/src/utils.rs +++ b/wayshot/src/utils.rs @@ -1,6 +1,11 @@ -use eyre::{ContextCompat, Result}; +use clap::ValueEnum; +use eyre::{bail, ContextCompat, Error, Result}; + use std::{ + fmt::Display, + path::PathBuf, process::exit, + str::FromStr, time::{SystemTime, UNIX_EPOCH}, }; @@ -47,18 +52,24 @@ pub fn parse_geometry(g: &str) -> Result { } /// Supported image encoding formats. -#[derive(Debug, Copy, Clone, PartialEq, Eq)] +#[derive(Debug, Copy, Clone, PartialEq, Eq, ValueEnum)] pub enum EncodingFormat { - /// Jpeg / jpg encoder. + /// JPG/JPEG encoder. Jpg, - /// Png encoder. + /// PNG encoder. Png, - /// Ppm encoder. + /// PPM encoder. Ppm, - /// Qoi encoder. + /// Qut encoder. Qoi, } +impl Default for EncodingFormat { + fn default() -> Self { + Self::Png + } +} + impl From for image::ImageOutputFormat { fn from(format: EncodingFormat) -> Self { match format { @@ -70,6 +81,33 @@ impl From for image::ImageOutputFormat { } } +impl TryFrom<&PathBuf> for EncodingFormat { + type Error = Error; + + fn try_from(value: &PathBuf) -> std::result::Result { + value + .extension() + .wrap_err_with(|| { + format!( + "no extension in {} to deduce encoding format", + value.display() + ) + }) + .and_then(|ext| { + ext.to_str().wrap_err_with(|| { + format!("extension in {} is not valid unicode", value.display()) + }) + }) + .and_then(|ext| ext.parse()) + } +} + +impl Display for EncodingFormat { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", Into::<&str>::into(*self)) + } +} + impl From for &str { fn from(format: EncodingFormat) -> Self { match format { @@ -81,7 +119,21 @@ impl From for &str { } } -pub fn get_default_file_name(extension: EncodingFormat) -> String { +impl FromStr for EncodingFormat { + type Err = Error; + + fn from_str(s: &str) -> std::result::Result { + Ok(match s { + "jpg" | "jpeg" => Self::Jpg, + "png" => Self::Png, + "ppm" => Self::Ppm, + "qoi" => Self::Qoi, + _ => bail!("unsupported extension '{s}'"), + }) + } +} + +pub fn get_default_file_name(extension: EncodingFormat) -> PathBuf { let time = match SystemTime::now().duration_since(UNIX_EPOCH) { Ok(n) => n.as_secs().to_string(), Err(_) => { @@ -90,5 +142,5 @@ pub fn get_default_file_name(extension: EncodingFormat) -> String { } }; - time + "-wayshot." + extension.into() + (time + "-wayshot." + extension.into()).into() } diff --git a/wayshot/src/wayshot.rs b/wayshot/src/wayshot.rs index 8e3bc8af..b5632caf 100644 --- a/wayshot/src/wayshot.rs +++ b/wayshot/src/wayshot.rs @@ -11,9 +11,7 @@ mod cli; mod utils; use dialoguer::{theme::ColorfulTheme, FuzzySelect}; -use tracing::Level; - -use crate::utils::EncodingFormat; +use utils::EncodingFormat; fn select_ouput(ouputs: &[T]) -> Option where @@ -32,41 +30,39 @@ where fn main() -> Result<()> { let cli = cli::Cli::parse(); - let level = if cli.debug { Level::TRACE } else { Level::INFO }; tracing_subscriber::fmt() - .with_max_level(level) + .with_max_level(cli.log_level) .with_writer(std::io::stderr) .init(); - let extension = if let Some(extension) = cli.extension { - let ext = extension.trim().to_lowercase(); - tracing::debug!("Using custom extension: {:#?}", ext); - - match ext.as_str() { - "jpeg" | "jpg" => EncodingFormat::Jpg, - "png" => EncodingFormat::Png, - "ppm" => EncodingFormat::Ppm, - "qoi" => EncodingFormat::Qoi, - _ => { - tracing::error!("Invalid extension provided.\nValid extensions:\n1) jpeg\n2) jpg\n3) png\n4) ppm\n5) qoi"); - exit(1); + let input_encoding = cli + .file + .as_ref() + .and_then(|pathbuf| pathbuf.try_into().ok()); + let requested_encoding = cli + .encoding + .or(input_encoding) + .unwrap_or(EncodingFormat::default()); + + if let Some(input_encoding) = input_encoding { + if input_encoding != requested_encoding { + tracing::warn!( + "The encoding requested '{requested_encoding}' does not match the output file's encoding '{input_encoding}'. Still using the requested encoding however.", + ); + } + } + + let file = match cli.file { + Some(pathbuf) => { + if pathbuf.to_string_lossy() == "-" { + None + } else { + Some(pathbuf) } } - } else { - EncodingFormat::Png + None => Some(utils::get_default_file_name(requested_encoding)), }; - let mut file_is_stdout: bool = false; - let mut file_path: Option = None; - - if cli.stdout { - file_is_stdout = true; - } else if let Some(filepath) = cli.file { - file_path = Some(filepath.trim().to_string()); - } else { - file_path = Some(utils::get_default_file_name(extension)); - } - let wayshot_conn = WayshotConnection::new()?; if cli.list_outputs { @@ -117,16 +113,16 @@ fn main() -> Result<()> { wayshot_conn.screenshot_all(cli.cursor)? }; - if file_is_stdout { + if let Some(file) = file { + image_buffer.save(file)?; + } else { let stdout = stdout(); let mut buffer = Cursor::new(Vec::new()); let mut writer = BufWriter::new(stdout.lock()); - image_buffer.write_to(&mut buffer, extension)?; + image_buffer.write_to(&mut buffer, requested_encoding)?; writer.write_all(buffer.get_ref())?; - } else { - image_buffer.save(file_path.unwrap())?; } Ok(()) From 7d6e20bdb2e6726fd0ca50778097af1174c6e0b8 Mon Sep 17 00:00:00 2001 From: Andreas Backx Date: Fri, 2 Feb 2024 19:01:37 +0000 Subject: [PATCH 5/6] Remove explicit panics and exits Let's bubble up errors instead of panicking and also not use `exit(...)` because it does not call destructors which might give issues. --- libwayshot/README.md | 4 +-- libwayshot/src/dispatch.rs | 23 +++++++++++------ libwayshot/src/error.rs | 7 +++++- libwayshot/src/lib.rs | 9 +++---- libwayshot/src/region.rs | 8 +++--- libwayshot/src/screencopy.rs | 25 ++++++++++-------- wayshot/Cargo.toml | 2 ++ wayshot/build.rs | 49 ++++++++++++++++++++++-------------- wayshot/src/cli.rs | 3 ++- wayshot/src/utils.rs | 14 ++++------- wayshot/src/wayshot.rs | 12 ++++----- 11 files changed, 89 insertions(+), 67 deletions(-) diff --git a/libwayshot/README.md b/libwayshot/README.md index a18983c7..a3bf2ab2 100644 --- a/libwayshot/README.md +++ b/libwayshot/README.md @@ -17,6 +17,6 @@ ```rust use libwayshot::WayshotConnection; -let wayshot_connection = WayshotConnection::new().unwrap(); -let image_buffer = wayshot_connection.screenshot_all().unwrap(); +let wayshot_connection = WayshotConnection::new()?; +let image_buffer = wayshot_connection.screenshot_all()?; ``` diff --git a/libwayshot/src/dispatch.rs b/libwayshot/src/dispatch.rs index f4eedb06..5352deb4 100644 --- a/libwayshot/src/dispatch.rs +++ b/libwayshot/src/dispatch.rs @@ -1,6 +1,5 @@ use std::{ collections::HashSet, - process::exit, sync::atomic::{AtomicBool, Ordering}, }; use wayland_client::{ @@ -86,11 +85,13 @@ impl Dispatch for OutputCaptureState { _: &Connection, _: &QueueHandle, ) { - let output: &mut OutputInfo = state - .outputs - .iter_mut() - .find(|x| x.wl_output == *wl_output) - .unwrap(); + let output: &mut OutputInfo = + if let Some(output) = state.outputs.iter_mut().find(|x| x.wl_output == *wl_output) { + output + } else { + tracing::error!("Received event for an output that is not registered: {event:#?}"); + return; + }; match event { wl_output::Event::Name { name } => { @@ -129,7 +130,14 @@ impl Dispatch for OutputCaptureState { _: &Connection, _: &QueueHandle, ) { - let output_info = state.outputs.get_mut(*index).unwrap(); + let output_info = if let Some(output_info) = state.outputs.get_mut(*index) { + output_info + } else { + tracing::error!( + "Received event for output index {index} that is not registered: {event:#?}" + ); + return; + }; match event { zxdg_output_v1::Event::LogicalPosition { x, y } => { @@ -189,7 +197,6 @@ impl Dispatch for CaptureFrameState { }) } else { tracing::debug!("Received Buffer event with unidentified format"); - exit(1); } } zwlr_screencopy_frame_v1::Event::Ready { .. } => { diff --git a/libwayshot/src/error.rs b/libwayshot/src/error.rs index ab0e1ec3..f3ad8523 100644 --- a/libwayshot/src/error.rs +++ b/libwayshot/src/error.rs @@ -1,7 +1,10 @@ use std::{io, result}; use thiserror::Error; -use wayland_client::{globals::GlobalError, ConnectError, DispatchError}; +use wayland_client::{ + globals::{BindError, GlobalError}, + ConnectError, DispatchError, +}; pub type Result = result::Result; @@ -17,6 +20,8 @@ pub enum Error { Io(#[from] io::Error), #[error("dispatch error: {0}")] Dispatch(#[from] DispatchError), + #[error("bind error: {0}")] + Bind(#[from] BindError), #[error("global error: {0}")] Global(#[from] GlobalError), #[error("connect error: {0}")] diff --git a/libwayshot/src/lib.rs b/libwayshot/src/lib.rs index 2128ed55..8198818c 100644 --- a/libwayshot/src/lib.rs +++ b/libwayshot/src/lib.rs @@ -15,7 +15,6 @@ use std::{ collections::HashSet, fs::File, os::fd::AsFd, - process::exit, sync::atomic::{AtomicBool, Ordering}, thread, }; @@ -68,8 +67,8 @@ pub mod reexport { /// # Example usage /// /// ``` -/// let wayshot_connection = WayshotConnection::new().unwrap(); -/// let image_buffer = wayshot_connection.screenshot_all().unwrap(); +/// let wayshot_connection = WayshotConnection::new()?; +/// let image_buffer = wayshot_connection.screenshot_all()?; /// ``` #[derive(Debug)] pub struct WayshotConnection { @@ -150,7 +149,7 @@ impl WayshotConnection { if state.outputs.is_empty() { tracing::error!("Compositor did not advertise any wl_output devices!"); - exit(1); + return Err(Error::NoOutputs); } tracing::trace!("Outputs detected: {:#?}", state.outputs); self.output_infos = state.outputs; @@ -280,7 +279,7 @@ impl WayshotConnection { let frame_bytes = frame_format.stride * frame_format.height; // Instantiate shm global. - let shm = self.globals.bind::(&qh, 1..=1, ()).unwrap(); + let shm = self.globals.bind::(&qh, 1..=1, ())?; let shm_pool = shm.create_pool(fd.as_fd(), frame_bytes as i32, &qh, ()); let buffer = shm_pool.create_buffer( 0, diff --git a/libwayshot/src/region.rs b/libwayshot/src/region.rs index 490e6082..27a35693 100644 --- a/libwayshot/src/region.rs +++ b/libwayshot/src/region.rs @@ -174,22 +174,22 @@ impl TryFrom<&Vec> for LogicalRegion { .iter() .map(|output| output.dimensions.x) .min() - .unwrap(); + .ok_or(Error::NoOutputs)?; let y1 = output_info .iter() .map(|output| output.dimensions.y) .min() - .unwrap(); + .ok_or(Error::NoOutputs)?; let x2 = output_info .iter() .map(|output| output.dimensions.x + output.dimensions.width) .max() - .unwrap(); + .ok_or(Error::NoOutputs)?; let y2 = output_info .iter() .map(|output| output.dimensions.y + output.dimensions.height) .max() - .unwrap(); + .ok_or(Error::NoOutputs)?; Ok(LogicalRegion { inner: Region { x: x1, diff --git a/libwayshot/src/screencopy.rs b/libwayshot/src/screencopy.rs index 2df15601..8d51b2ce 100644 --- a/libwayshot/src/screencopy.rs +++ b/libwayshot/src/screencopy.rs @@ -1,5 +1,5 @@ use std::{ - ffi::CStr, + ffi::CString, os::fd::{AsRawFd, IntoRawFd, OwnedFd}, time::{SystemTime, UNIX_EPOCH}, }; @@ -82,6 +82,16 @@ impl TryFrom<&FrameCopy> for DynamicImage { } } +fn get_mem_file_handle() -> String { + format!( + "/libwayshot-{}", + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|time| time.subsec_nanos().to_string()) + .unwrap_or("unknown".into()) + ) +} + /// Return a RawFd to a shm file. We use memfd create on linux and shm_open for BSD support. /// You don't need to mess around with this function, it is only used by /// capture_output_frame. @@ -91,7 +101,7 @@ pub fn create_shm_fd() -> std::io::Result { loop { // Create a file that closes on succesful execution and seal it's operations. match memfd::memfd_create( - CStr::from_bytes_with_nul(b"libwayshot\0").unwrap(), + CString::new("libwayshot")?.as_c_str(), memfd::MemFdCreateFlag::MFD_CLOEXEC | memfd::MemFdCreateFlag::MFD_ALLOW_SEALING, ) { Ok(fd) => { @@ -113,11 +123,7 @@ pub fn create_shm_fd() -> std::io::Result { } // Fallback to using shm_open. - let sys_time = SystemTime::now(); - let mut mem_file_handle = format!( - "/libwayshot-{}", - sys_time.duration_since(UNIX_EPOCH).unwrap().subsec_nanos() - ); + let mut mem_file_handle = get_mem_file_handle(); loop { match mman::shm_open( // O_CREAT = Create file if does not exist. @@ -142,10 +148,7 @@ pub fn create_shm_fd() -> std::io::Result { }, Err(nix::errno::Errno::EEXIST) => { // If a file with that handle exists then change the handle - mem_file_handle = format!( - "/libwayshot-{}", - sys_time.duration_since(UNIX_EPOCH).unwrap().subsec_nanos() - ); + mem_file_handle = get_mem_file_handle(); continue; } Err(nix::errno::Errno::EINTR) => continue, diff --git a/wayshot/Cargo.toml b/wayshot/Cargo.toml index 344654f9..4594d0b4 100644 --- a/wayshot/Cargo.toml +++ b/wayshot/Cargo.toml @@ -12,6 +12,8 @@ repository.workspace = true [build-dependencies] flate2 = "1.0.27" +eyre = "0.6.8" + [dependencies] tracing.workspace = true diff --git a/wayshot/build.rs b/wayshot/build.rs index 14152acb..37104611 100644 --- a/wayshot/build.rs +++ b/wayshot/build.rs @@ -1,13 +1,14 @@ extern crate flate2; +use eyre::{ContextCompat, Result}; use flate2::{write::GzEncoder, Compression}; use std::{ fs::{read_dir, File, OpenOptions}, io::{copy, BufReader, ErrorKind}, path::Path, - process::{exit, Command, Stdio}, + process::{Command, Stdio}, }; -fn main() { +fn main() -> Result<()> { if let Err(e) = Command::new("scdoc") .stdin(Stdio::null()) .stdout(Stdio::null()) @@ -15,50 +16,60 @@ fn main() { .spawn() { if let ErrorKind::NotFound = e.kind() { - exit(0); + return Ok(()); } } // We just append "out" so it's easy to find all the scdoc output later in line 38. - let man_pages: Vec<(String, String)> = read_and_replace_by_ext("./docs", ".scd", ".out"); + let man_pages: Vec<(String, String)> = read_and_replace_by_ext("./docs", ".scd", ".out")?; for man_page in man_pages { let output = OpenOptions::new() .write(true) .create(true) - .open(Path::new(&man_page.1)) - .unwrap(); + .open(Path::new(&man_page.1))?; _ = Command::new("scdoc") - .stdin(Stdio::from(File::open(man_page.0).unwrap())) + .stdin(Stdio::from(File::open(man_page.0)?)) .stdout(output) .spawn(); } // Gzipping the man pages let scdoc_output_files: Vec<(String, String)> = - read_and_replace_by_ext("./docs", ".out", ".gz"); + read_and_replace_by_ext("./docs", ".out", ".gz")?; for scdoc_output in scdoc_output_files { - let mut input = BufReader::new(File::open(scdoc_output.0).unwrap()); + let mut input = BufReader::new(File::open(scdoc_output.0)?); let output = OpenOptions::new() .write(true) .create(true) - .open(Path::new(&scdoc_output.1)) - .unwrap(); + .open(Path::new(&scdoc_output.1))?; let mut encoder = GzEncoder::new(output, Compression::default()); - copy(&mut input, &mut encoder).unwrap(); - encoder.finish().unwrap(); + copy(&mut input, &mut encoder)?; + encoder.finish()?; } + + Ok(()) } -fn read_and_replace_by_ext(path: &str, search: &str, replace: &str) -> Vec<(String, String)> { +fn read_and_replace_by_ext( + path: &str, + search: &str, + replace: &str, +) -> Result> { let mut files: Vec<(String, String)> = Vec::new(); - for path in read_dir(path).unwrap() { - let path = path.unwrap(); - if path.file_type().unwrap().is_dir() { + for path in read_dir(path)? { + let path = path?; + if path.file_type()?.is_dir() { continue; } if let Some(file_name) = path.path().to_str() { - if *path.path().extension().unwrap().to_str().unwrap() != search[1..] { + if *path + .path() + .extension() + .wrap_err_with(|| format!("no extension found for {}", path.path().display()))? + .to_string_lossy() + != search[1..] + { continue; } @@ -66,5 +77,5 @@ fn read_and_replace_by_ext(path: &str, search: &str, replace: &str) -> Vec<(Stri files.push((file_name.to_string(), file)); } } - files + Ok(files) } diff --git a/wayshot/src/cli.rs b/wayshot/src/cli.rs index a076af5c..418ac8d8 100644 --- a/wayshot/src/cli.rs +++ b/wayshot/src/cli.rs @@ -3,6 +3,7 @@ use std::path::PathBuf; use clap::arg; use clap::Parser; +use eyre::WrapErr; use crate::utils::EncodingFormat; use clap::builder::TypedValueParser; @@ -15,7 +16,7 @@ pub struct Cli { pub file: Option, /// Log level to be used for printing to stderr - #[arg(long, default_value = "info", value_parser = clap::builder::PossibleValuesParser::new(["trace", "debug", "info", "warn", "error"]).map(|s| -> tracing::Level{ s.parse().unwrap()}))] + #[arg(long, default_value = "info", value_parser = clap::builder::PossibleValuesParser::new(["trace", "debug", "info", "warn", "error"]).map(|s| -> tracing::Level{ s.parse().wrap_err_with(|| format!("Failed to parse log level: {}", s)).unwrap()}))] pub log_level: tracing::Level, /// Arguments to call slurp with for selecting a region diff --git a/wayshot/src/utils.rs b/wayshot/src/utils.rs index 1e394452..bb9a7a96 100644 --- a/wayshot/src/utils.rs +++ b/wayshot/src/utils.rs @@ -4,7 +4,6 @@ use eyre::{bail, ContextCompat, Error, Result}; use std::{ fmt::Display, path::PathBuf, - process::exit, str::FromStr, time::{SystemTime, UNIX_EPOCH}, }; @@ -134,13 +133,10 @@ impl FromStr for EncodingFormat { } pub fn get_default_file_name(extension: EncodingFormat) -> PathBuf { - let time = match SystemTime::now().duration_since(UNIX_EPOCH) { - Ok(n) => n.as_secs().to_string(), - Err(_) => { - tracing::error!("SystemTime before UNIX EPOCH!"); - exit(1); - } - }; + let time = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|time| time.as_secs().to_string()) + .unwrap_or("unknown".into()); - (time + "-wayshot." + extension.into()).into() + format!("{time}-wayshot.{extension}").into() } diff --git a/wayshot/src/wayshot.rs b/wayshot/src/wayshot.rs index b5632caf..37152a79 100644 --- a/wayshot/src/wayshot.rs +++ b/wayshot/src/wayshot.rs @@ -1,10 +1,10 @@ use std::{ io::{stdout, BufWriter, Cursor, Write}, - process::{exit, Command}, + process::Command, }; use clap::Parser; -use eyre::Result; +use eyre::{bail, Result}; use libwayshot::{region::LogicalRegion, WayshotConnection}; mod cli; @@ -70,7 +70,7 @@ fn main() -> Result<()> { for output in valid_outputs { tracing::info!("{:#?}", output.name); } - exit(1); + return Ok(()); } let image_buffer = if let Some(slurp_region) = cli.slurp { @@ -94,8 +94,7 @@ fn main() -> Result<()> { if let Some(output) = outputs.iter().find(|output| output.name == output_name) { wayshot_conn.screenshot_single_output(output, cli.cursor)? } else { - tracing::error!("No output found!\n"); - exit(1); + bail!("No output found!"); } } else if cli.choose_output { let outputs = wayshot_conn.get_all_outputs(); @@ -106,8 +105,7 @@ fn main() -> Result<()> { if let Some(index) = select_ouput(&output_names) { wayshot_conn.screenshot_single_output(&outputs[index], cli.cursor)? } else { - tracing::error!("No output found!\n"); - exit(1); + bail!("No output found!"); } } else { wayshot_conn.screenshot_all(cli.cursor)? From 044910d6667582a4f53cdccf40c4a15f2bdd916c Mon Sep 17 00:00:00 2001 From: Andreas Backx Date: Fri, 2 Feb 2024 19:19:00 +0000 Subject: [PATCH 6/6] Make region, sizing, and positioning data structs reusable We use different structs/fields to store the same three types of data everywhere: - Position (x,y) - Size (width,height) - Region(position, size) This makes it so they all reuse the same information. --- libwayshot/src/dispatch.rs | 22 +++--- libwayshot/src/lib.rs | 92 +++++++++++++---------- libwayshot/src/output.rs | 28 +++---- libwayshot/src/region.rs | 140 ++++++++++++++++++----------------- libwayshot/src/screencopy.rs | 28 +++++-- wayshot/src/utils.rs | 23 +++--- 6 files changed, 183 insertions(+), 150 deletions(-) diff --git a/libwayshot/src/dispatch.rs b/libwayshot/src/dispatch.rs index 5352deb4..574fc813 100644 --- a/libwayshot/src/dispatch.rs +++ b/libwayshot/src/dispatch.rs @@ -26,7 +26,8 @@ use wayland_protocols_wlr::screencopy::v1::client::{ }; use crate::{ - output::{OutputInfo, OutputPositioning, WlOutputMode}, + output::OutputInfo, + region::{Position, Region, Size}, screencopy::FrameFormat, }; @@ -64,8 +65,7 @@ impl Dispatch for OutputCaptureState { description: String::new(), transform: wl_output::Transform::Normal, scale: 1, - dimensions: OutputPositioning::default(), - mode: WlOutputMode::default(), + region: Region::default(), }); } else { tracing::error!("Ignoring a wl_output with version < 4."); @@ -100,9 +100,7 @@ impl Dispatch for OutputCaptureState { wl_output::Event::Description { description } => { output.description = description; } - wl_output::Event::Mode { width, height, .. } => { - output.mode = WlOutputMode { width, height }; - } + wl_output::Event::Mode { .. } => {} wl_output::Event::Geometry { transform: WEnum::Value(transform), .. @@ -141,12 +139,13 @@ impl Dispatch for OutputCaptureState { match event { zxdg_output_v1::Event::LogicalPosition { x, y } => { - output_info.dimensions.x = x; - output_info.dimensions.y = y; + output_info.region.position = Position { x, y }; } zxdg_output_v1::Event::LogicalSize { width, height } => { - output_info.dimensions.width = width; - output_info.dimensions.height = height; + output_info.region.size = Size { + width: width as u32, + height: height as u32, + }; } zxdg_output_v1::Event::Done => {} zxdg_output_v1::Event::Name { .. } => {} @@ -191,8 +190,7 @@ impl Dispatch for CaptureFrameState { if let Value(f) = format { frame.formats.push(FrameFormat { format: f, - width, - height, + size: Size { width, height }, stride, }) } else { diff --git a/libwayshot/src/lib.rs b/libwayshot/src/lib.rs index 8198818c..beda48fe 100644 --- a/libwayshot/src/lib.rs +++ b/libwayshot/src/lib.rs @@ -29,7 +29,7 @@ use wayland_client::{ globals::{registry_queue_init, GlobalList}, protocol::{ wl_compositor::WlCompositor, - wl_output::WlOutput, + wl_output::{Transform, WlOutput}, wl_shm::{self, WlShm}, }, Connection, EventQueue, @@ -52,7 +52,7 @@ use crate::{ convert::create_converter, dispatch::{CaptureFrameState, FrameState, OutputCaptureState, WayshotState}, output::OutputInfo, - region::LogicalRegion, + region::{LogicalRegion, Region, Size}, screencopy::{create_shm_fd, FrameCopy, FrameFormat}, }; @@ -214,10 +214,10 @@ impl WayshotConnection { screencopy_manager.capture_output_region( cursor_overlay, output, - embedded_region.inner.x, - embedded_region.inner.y, - embedded_region.inner.width, - embedded_region.inner.height, + embedded_region.inner.position.x, + embedded_region.inner.position.y, + embedded_region.inner.size.width as i32, + embedded_region.inner.size.height as i32, &qh, (), ) @@ -275,16 +275,21 @@ impl WayshotConnection { // Connecting to wayland environment. let qh = event_queue.handle(); - // Bytes of data in the frame = stride * height. - let frame_bytes = frame_format.stride * frame_format.height; - // Instantiate shm global. let shm = self.globals.bind::(&qh, 1..=1, ())?; - let shm_pool = shm.create_pool(fd.as_fd(), frame_bytes as i32, &qh, ()); + let shm_pool = shm.create_pool( + fd.as_fd(), + frame_format + .byte_size() + .try_into() + .map_err(|_| Error::BufferTooSmall)?, + &qh, + (), + ); let buffer = shm_pool.create_buffer( 0, - frame_format.width as i32, - frame_format.height as i32, + frame_format.size.width as i32, + frame_format.size.height as i32, frame_format.stride as i32, frame_format.format, &qh, @@ -322,9 +327,7 @@ impl WayshotConnection { let (state, event_queue, frame, frame_format) = self.capture_output_frame_get_state(cursor_overlay as i32, output, capture_region)?; - // Bytes of data in the frame = stride * height. - let frame_bytes = frame_format.stride * frame_format.height; - file.set_len(frame_bytes as u64)?; + file.set_len(frame_format.byte_size())?; let frame_guard = self.capture_output_frame_inner(state, event_queue, frame, frame_format, file)?; @@ -333,7 +336,7 @@ impl WayshotConnection { } /// Get a FrameCopy instance with screenshot pixel data for any wl_output object. - #[tracing::instrument(skip_all, fields(output = output_info.name, region = capture_region.map(|r| format!("{:}", r))))] + #[tracing::instrument(skip_all, fields(output = format!("{output_info}"), region = capture_region.map(|r| format!("{:}", r))))] fn capture_frame_copy( &self, cursor_overlay: bool, @@ -361,22 +364,30 @@ impl WayshotConnection { tracing::error!("You can send a feature request for the above format to the mailing list for wayshot over at https://sr.ht/~shinyzenith/wayshot."); return Err(Error::NoSupportedBufferFormat); }; + let rotated_size = match output_info.transform { + Transform::_90 | Transform::_270 | Transform::Flipped90 | Transform::Flipped270 => { + Size { + width: frame_format.size.height, + height: frame_format.size.width, + } + } + _ => frame_format.size, + }; let frame_copy = FrameCopy { frame_format, frame_color_type, frame_mmap, transform: output_info.transform, - position: capture_region - .map(|capture_region| { - let logical_region = capture_region.logical(); - (logical_region.inner.x as i64, logical_region.inner.y as i64) - }) - .unwrap_or_else(|| { - ( - output_info.dimensions.x as i64, - output_info.dimensions.y as i64, - ) - }), + region: LogicalRegion { + inner: Region { + position: capture_region + .map(|capture_region: EmbeddedRegion| { + capture_region.logical().inner.position + }) + .unwrap_or_else(|| output_info.region.position), + size: rotated_size, + }, + }, }; tracing::debug!("Created frame copy: {:#?}", frame_copy); Ok((frame_copy, frame_guard)) @@ -451,7 +462,7 @@ impl WayshotConnection { tracing::span!( tracing::Level::DEBUG, "overlay_frames::surface", - output = output_info.name.as_str() + output = format!("{output_info}") ) .in_scope(|| -> Result<()> { let surface = compositor.create_surface(&qh, ()); @@ -468,8 +479,8 @@ impl WayshotConnection { layer_surface.set_exclusive_zone(-1); layer_surface.set_anchor(Anchor::Top | Anchor::Left); layer_surface.set_size( - frame_copy.frame_format.width, - frame_copy.frame_format.height, + frame_copy.frame_format.size.width, + frame_copy.frame_format.size.height, ); debug!("Committing surface creation changes."); @@ -515,8 +526,8 @@ impl WayshotConnection { tracing::Level::DEBUG, "filter_map", output = format!( - "{name} at {region}", - name = output_info.name, + "{output_info} at {region}", + output_info = format!("{output_info}"), region = LogicalRegion::from(output_info), ), capture_region = format!("{}", capture_region), @@ -561,8 +572,8 @@ impl WayshotConnection { image_util::rotate_image_buffer( image, frame_copy.transform, - frame_copy.frame_format.width, - frame_copy.frame_format.height, + frame_copy.frame_format.size.width, + frame_copy.frame_format.size.height, ), frame_copy, )) @@ -580,23 +591,24 @@ impl WayshotConnection { // Default to a transparent image. let composite_image = composite_image.unwrap_or_else(|| { Ok(DynamicImage::new_rgba8( - capture_region.inner.width as u32, - capture_region.inner.height as u32, + capture_region.inner.size.width, + capture_region.inner.size.height, )) }); Some(|| -> Result<_> { let mut composite_image = composite_image?; let (image, frame_copy) = image?; - let frame_copy_region = LogicalRegion::from(&frame_copy); let (x, y) = ( - frame_copy_region.inner.x as i64 - capture_region.inner.x as i64, - frame_copy_region.inner.y as i64 - capture_region.inner.y as i64, + frame_copy.region.inner.position.x as i64 + - capture_region.inner.position.x as i64, + frame_copy.region.inner.position.y as i64 + - capture_region.inner.position.y as i64, ); tracing::span!( tracing::Level::DEBUG, "replace", - frame_copy_region = format!("{}", frame_copy_region), + frame_copy_region = format!("{}", frame_copy.region), capture_region = format!("{}", capture_region), x = x, y = y, diff --git a/libwayshot/src/output.rs b/libwayshot/src/output.rs index 986f74df..5449f4f7 100644 --- a/libwayshot/src/output.rs +++ b/libwayshot/src/output.rs @@ -1,5 +1,9 @@ +use std::fmt::Display; + use wayland_client::protocol::{wl_output, wl_output::WlOutput}; +use crate::region::Region; + /// Represents an accessible wayland output. /// /// Do not instantiate, instead use [`crate::WayshotConnection::get_all_outputs`]. @@ -10,20 +14,16 @@ pub struct OutputInfo { pub description: String, pub transform: wl_output::Transform, pub scale: i32, - pub dimensions: OutputPositioning, - pub mode: WlOutputMode, -} - -#[derive(Default, Debug, Clone, PartialEq, Eq, Hash)] -pub struct WlOutputMode { - pub width: i32, - pub height: i32, + pub region: Region, } -#[derive(Default, Debug, Clone, PartialEq, Eq, Hash)] -pub struct OutputPositioning { - pub x: i32, - pub y: i32, - pub width: i32, - pub height: i32, +impl Display for OutputInfo { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "{name} ({description})", + name = self.name, + description = self.description + ) + } } diff --git a/libwayshot/src/region.rs b/libwayshot/src/region.rs index 27a35693..bbd81d39 100644 --- a/libwayshot/src/region.rs +++ b/libwayshot/src/region.rs @@ -1,10 +1,7 @@ use std::cmp; -use wayland_client::protocol::wl_output::Transform; - use crate::error::{Error, Result}; use crate::output::OutputInfo; -use crate::screencopy::FrameCopy; /// Ways to say how a region for a screenshot should be captured. pub enum RegionCapturer { @@ -58,16 +55,28 @@ pub struct EmbeddedRegion { /// /// Use `LogicalRegion` or `EmbeddedRegion` instead as they convey the /// coordinate system used. -#[derive(Debug, Copy, Clone)] +#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash, Default)] pub struct Region { - /// X coordinate of the area to capture. + /// Position of the region. + pub position: Position, + /// Size of the region. + pub size: Size, +} + +#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash, Default)] +pub struct Position { + /// X coordinate. pub x: i32, - /// y coordinate of the area to capture. + /// Y coordinate. pub y: i32, - /// Width of the capture area. - pub width: i32, - /// Height of the capture area. - pub height: i32, +} + +#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash, Default)] +pub struct Size { + /// Width. + pub width: u32, + /// Height. + pub height: u32, } impl EmbeddedRegion { @@ -78,30 +87,36 @@ impl EmbeddedRegion { /// See `EmbeddedRegion` for an example ASCII visualisation. #[tracing::instrument(ret, level = "debug")] pub fn new(viewport: LogicalRegion, relative_to: LogicalRegion) -> Option { - let x_relative: i32 = viewport.inner.x - relative_to.inner.x; - let y_relative = viewport.inner.y - relative_to.inner.y; + let x_relative: i32 = viewport.inner.position.x - relative_to.inner.position.x; + let y_relative = viewport.inner.position.y - relative_to.inner.position.y; let x1 = cmp::max(x_relative, 0); - let x2 = cmp::min(x_relative + viewport.inner.width, relative_to.inner.width); - let width = x2 - x1; - if width <= 0 { + let x2 = cmp::min( + x_relative + viewport.inner.size.width as i32, + relative_to.inner.size.width as i32, + ); + let width = if let Ok(width) = (x2 - x1).try_into() { + width + } else { return None; - } + }; let y1 = cmp::max(y_relative, 0); - let y2 = cmp::min(y_relative + viewport.inner.height, relative_to.inner.height); - let height = y2 - y1; - if height <= 0 { + let y2 = cmp::min( + y_relative + viewport.inner.size.height as i32, + relative_to.inner.size.height as i32, + ); + let height = if let Ok(height) = (y2 - y1).try_into() { + height + } else { return None; - } + }; Some(Self { relative_to: relative_to, inner: Region { - x: x1, - y: y1, - width, - height, + position: Position { x: x1, y: y1 }, + size: Size { width, height }, }, }) } @@ -114,10 +129,11 @@ impl EmbeddedRegion { pub fn logical(&self) -> LogicalRegion { LogicalRegion { inner: Region { - x: self.relative_to.inner.x as i32 + self.inner.x, - y: self.relative_to.inner.y as i32 + self.inner.y, - width: self.inner.width, - height: self.inner.height, + position: Position { + x: self.relative_to.inner.position.x + self.inner.position.x, + y: self.relative_to.inner.position.y + self.inner.position.y, + }, + size: self.inner.size, }, } } @@ -134,19 +150,34 @@ impl std::fmt::Display for EmbeddedRegion { } } -impl std::fmt::Display for Region { +impl std::fmt::Display for Position { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "({x}, {y})", x = self.x, y = self.y,) + } +} + +impl std::fmt::Display for Size { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!( f, - "({x}, {y}) ({width}x{height})", - x = self.x, - y = self.y, + "({width}x{height})", width = self.width, height = self.height, ) } } +impl std::fmt::Display for Region { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "({position}) ({size})", + position = self.position, + size = self.size, + ) + } +} + impl std::fmt::Display for LogicalRegion { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!(f, "{inner}", inner = self.inner) @@ -156,12 +187,7 @@ impl std::fmt::Display for LogicalRegion { impl From<&OutputInfo> for LogicalRegion { fn from(output_info: &OutputInfo) -> Self { LogicalRegion { - inner: Region { - x: output_info.dimensions.x, - y: output_info.dimensions.y, - width: output_info.dimensions.width, - height: output_info.dimensions.height, - }, + inner: output_info.region, } } } @@ -172,52 +198,32 @@ impl TryFrom<&Vec> for LogicalRegion { fn try_from(output_info: &Vec) -> std::result::Result { let x1 = output_info .iter() - .map(|output| output.dimensions.x) + .map(|output| output.region.position.x) .min() .ok_or(Error::NoOutputs)?; let y1 = output_info .iter() - .map(|output| output.dimensions.y) + .map(|output| output.region.position.y) .min() .ok_or(Error::NoOutputs)?; let x2 = output_info .iter() - .map(|output| output.dimensions.x + output.dimensions.width) + .map(|output| output.region.position.x + output.region.size.width as i32) .max() .ok_or(Error::NoOutputs)?; let y2 = output_info .iter() - .map(|output| output.dimensions.y + output.dimensions.height) + .map(|output| output.region.position.y + output.region.size.height as i32) .max() .ok_or(Error::NoOutputs)?; Ok(LogicalRegion { inner: Region { - x: x1, - y: y1, - width: x2 - x1, - height: y2 - y1, + position: Position { x: x1, y: y1 }, + size: Size { + width: (x2 - x1) as u32, + height: (y2 - y1) as u32, + }, }, }) } } - -impl From<&FrameCopy> for LogicalRegion { - fn from(frame_copy: &FrameCopy) -> Self { - let (width, height) = ( - frame_copy.frame_format.width as i32, - frame_copy.frame_format.height as i32, - ); - let is_portait = match frame_copy.transform { - Transform::_90 | Transform::_270 | Transform::Flipped90 | Transform::Flipped270 => true, - _ => false, - }; - LogicalRegion { - inner: Region { - x: frame_copy.position.0 as i32, - y: frame_copy.position.1 as i32, - width: if is_portait { height } else { width }, - height: if is_portait { width } else { height }, - }, - } - } -} diff --git a/libwayshot/src/screencopy.rs b/libwayshot/src/screencopy.rs index 8d51b2ce..734b6bfd 100644 --- a/libwayshot/src/screencopy.rs +++ b/libwayshot/src/screencopy.rs @@ -15,7 +15,10 @@ use wayland_client::protocol::{ wl_buffer::WlBuffer, wl_output, wl_shm::Format, wl_shm_pool::WlShmPool, }; -use crate::{Error, Result}; +use crate::{ + region::{LogicalRegion, Size}, + Error, Result, +}; pub struct FrameGuard { pub buffer: WlBuffer, @@ -36,12 +39,20 @@ impl Drop for FrameGuard { #[derive(Debug, Copy, Clone, PartialEq)] pub struct FrameFormat { pub format: Format, - pub width: u32, - pub height: u32, + /// Size of the frame in pixels. This will always be in "landscape" so a + /// portrait 1080x1920 frame will be 1920x1080 and will need to be rotated! + pub size: Size, /// Stride is the number of bytes between the start of a row and the start of the next row. pub stride: u32, } +impl FrameFormat { + /// Returns the size of the frame in bytes, which is the stride * height. + pub fn byte_size(&self) -> u64 { + self.stride as u64 * self.size.height as u64 + } +} + #[tracing::instrument(skip(frame_mmap))] fn create_image_buffer

( frame_format: &FrameFormat, @@ -51,8 +62,12 @@ where P: Pixel, { tracing::debug!("Creating image buffer"); - ImageBuffer::from_vec(frame_format.width, frame_format.height, frame_mmap.to_vec()) - .ok_or(Error::BufferTooSmall) + ImageBuffer::from_vec( + frame_format.size.width, + frame_format.size.height, + frame_mmap.to_vec(), + ) + .ok_or(Error::BufferTooSmall) } /// The copied frame comprising of the FrameFormat, ColorType (Rgba8), and a memory backed shm @@ -63,7 +78,8 @@ pub struct FrameCopy { pub frame_color_type: ColorType, pub frame_mmap: MmapMut, pub transform: wl_output::Transform, - pub position: (i64, i64), + /// Logical region with the transform already applied. + pub region: LogicalRegion, } impl TryFrom<&FrameCopy> for DynamicImage { diff --git a/wayshot/src/utils.rs b/wayshot/src/utils.rs index bb9a7a96..117c1f0c 100644 --- a/wayshot/src/utils.rs +++ b/wayshot/src/utils.rs @@ -8,14 +8,14 @@ use std::{ time::{SystemTime, UNIX_EPOCH}, }; -use libwayshot::region::{LogicalRegion, Region}; +use libwayshot::region::{LogicalRegion, Position, Region, Size}; pub fn parse_geometry(g: &str) -> Result { let tail = g.trim(); let x_coordinate: i32; let y_coordinate: i32; - let width: i32; - let height: i32; + let width: u32; + let height: u32; let validation_error = "Invalid geometry provided.\nValid geometries:\n1) %d,%d %dx%d\n2) %d %d %d %d"; @@ -27,8 +27,8 @@ pub fn parse_geometry(g: &str) -> Result { let (head, tail) = tail.split_once(' ').wrap_err(validation_error)?; y_coordinate = head.parse::()?; let (head, tail) = tail.split_once('x').wrap_err(validation_error)?; - width = head.parse::()?; - height = tail.parse::()?; + width = head.parse::()?; + height = tail.parse::()?; } else { // this accepts: "%d %d %d %d" let (head, tail) = tail.split_once(' ').wrap_err(validation_error)?; @@ -36,16 +36,17 @@ pub fn parse_geometry(g: &str) -> Result { let (head, tail) = tail.split_once(' ').wrap_err(validation_error)?; y_coordinate = head.parse::()?; let (head, tail) = tail.split_once(' ').wrap_err(validation_error)?; - width = head.parse::()?; - height = tail.parse::()?; + width = head.parse::()?; + height = tail.parse::()?; } Ok(LogicalRegion { inner: Region { - x: x_coordinate, - y: y_coordinate, - width, - height, + position: Position { + x: x_coordinate, + y: y_coordinate, + }, + size: Size { width, height }, }, }) }