[libwayshot] Mutli-monitor handling built into libwayshot

Signed-off-by: c-h-johnson
Signed-off-by: decodetalkers
Signed-off-by: Shinyzenith <aakashsensharma@gmail.com>

libwayshot/Cargo.toml

@@ -9,10 +9,11 @@ version = "0.1.2"
 edition = "2021"
 
 [dependencies]
-image = { version = "0.24", default-features = false, features = ["jpeg", "png", "pnm", "qoi"] }
+image = { version = "0.24", default-features = false }
 log = "0.4.19"
 memmap2 = "0.7.1"
 nix = "0.26.2"
+thiserror = "1"
 wayland-client = "0.30.2"
 wayland-protocols = { version = "0.30.0", features=["client", "unstable"] }
 wayland-protocols-wlr = { version = "0.1.0", features = ["client"] }

libwayshot/README.md

@@ -8,3 +8,15 @@
   </p>
 </p>
 
+# `libwayshot`
+
+`libwayshot` is a convenient wrapper over the wlroots screenshot protocol that provides a simple API to take screenshots with.
+
+# Example usage
+
+```rust
+use libwayshot::WayshotConnection;
+
+let wayshot_connection = WayshotConnection::new().unwrap();
+let image_buffer = wayshot_connection.screenshot_all().unwrap();
+```

libwayshot/src/dispatch.rs

@@ -0,0 +1,223 @@
+use std::{
+    process::exit,
+    sync::atomic::{AtomicBool, Ordering},
+};
+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,
+    },
+    Connection, Dispatch, QueueHandle, WEnum,
+    WEnum::Value,
+};
+use wayland_protocols::xdg::xdg_output::zv1::client::{
+    zxdg_output_manager_v1::ZxdgOutputManagerV1, zxdg_output_v1, zxdg_output_v1::ZxdgOutputV1,
+};
+use wayland_protocols_wlr::screencopy::v1::client::{
+    zwlr_screencopy_frame_v1, zwlr_screencopy_frame_v1::ZwlrScreencopyFrameV1,
+    zwlr_screencopy_manager_v1::ZwlrScreencopyManagerV1,
+};
+
+use crate::{
+    output::{OutputInfo, OutputPositioning},
+    screencopy::FrameFormat,
+};
+
+pub struct OutputCaptureState {
+    pub outputs: Vec<OutputInfo>,
+}
+
+impl Dispatch<WlRegistry, ()> for OutputCaptureState {
+    fn event(
+        state: &mut Self,
+        wl_registry: &WlRegistry,
+        event: wl_registry::Event,
+        _: &(),
+        _: &Connection,
+        qh: &QueueHandle<Self>,
+    ) {
+        /* > The name event is sent after binding the output object. This event
+         * is only sent once per output object, and the name does not change
+         * over the lifetime of the wl_output global. */
+
+        if let wl_registry::Event::Global {
+            name,
+            interface,
+            version,
+        } = event
+        {
+            if interface == "wl_output" {
+                if version >= 4 {
+                    let output = wl_registry.bind::<wl_output::WlOutput, _, _>(name, 4, qh, ());
+                    state.outputs.push(OutputInfo {
+                        wl_output: output,
+                        name: "".to_string(),
+                        transform: wl_output::Transform::Normal,
+                        dimensions: OutputPositioning {
+                            x: 0,
+                            y: 0,
+                            width: 0,
+                            height: 0,
+                        },
+                    });
+                } else {
+                    log::error!("Ignoring a wl_output with version < 4.");
+                }
+            }
+        }
+    }
+}
+
+impl Dispatch<WlOutput, ()> for OutputCaptureState {
+    fn event(
+        state: &mut Self,
+        wl_output: &WlOutput,
+        event: wl_output::Event,
+        _: &(),
+        _: &Connection,
+        _: &QueueHandle<Self>,
+    ) {
+        let output: &mut OutputInfo = state
+            .outputs
+            .iter_mut()
+            .find(|x| x.wl_output == *wl_output)
+            .unwrap();
+
+        match event {
+            wl_output::Event::Name { name } => {
+                output.name = name;
+            }
+            wl_output::Event::Geometry {
+                transform: WEnum::Value(transform),
+                ..
+            } => {
+                output.transform = transform;
+            }
+            _ => (),
+        }
+    }
+}
+
+delegate_noop!(OutputCaptureState: ignore ZxdgOutputManagerV1);
+
+impl Dispatch<ZxdgOutputV1, usize> for OutputCaptureState {
+    fn event(
+        state: &mut Self,
+        _: &ZxdgOutputV1,
+        event: zxdg_output_v1::Event,
+        index: &usize,
+        _: &Connection,
+        _: &QueueHandle<Self>,
+    ) {
+        let output_info = state.outputs.get_mut(*index).unwrap();
+
+        match event {
+            zxdg_output_v1::Event::LogicalPosition { x, y } => {
+                output_info.dimensions.x = x;
+                output_info.dimensions.y = y;
+                log::debug!("Logical position event fired!");
+            }
+            zxdg_output_v1::Event::LogicalSize { width, height } => {
+                output_info.dimensions.width = width;
+                output_info.dimensions.height = height;
+                log::debug!("Logical size event fired!");
+            }
+            _ => {}
+        };
+    }
+}
+
+/// State of the frame after attemting 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`.
+    Failed,
+    /// Compositor sent a Ready event on calling `frame.copy`.
+    Finished,
+}
+
+pub struct CaptureFrameState {
+    pub formats: Vec<FrameFormat>,
+    pub state: Option<FrameState>,
+    pub buffer_done: AtomicBool,
+}
+
+impl Dispatch<ZwlrScreencopyFrameV1, ()> for CaptureFrameState {
+    fn event(
+        frame: &mut Self,
+        _: &ZwlrScreencopyFrameV1,
+        event: zwlr_screencopy_frame_v1::Event,
+        _: &(),
+        _: &Connection,
+        _: &QueueHandle<Self>,
+    ) {
+        match event {
+            zwlr_screencopy_frame_v1::Event::Buffer {
+                format,
+                width,
+                height,
+                stride,
+            } => {
+                log::debug!("Received Buffer event");
+                if let Value(f) = format {
+                    frame.formats.push(FrameFormat {
+                        format: f,
+                        width,
+                        height,
+                        stride,
+                    })
+                } else {
+                    log::debug!("Received Buffer event with unidentified format");
+                    exit(1);
+                }
+            }
+            zwlr_screencopy_frame_v1::Event::Flags { .. } => {
+                log::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.
+                log::debug!("Received Ready event");
+                frame.state.replace(FrameState::Finished);
+            }
+            zwlr_screencopy_frame_v1::Event::Failed => {
+                log::debug!("Received Failed event");
+                frame.state.replace(FrameState::Failed);
+            }
+            zwlr_screencopy_frame_v1::Event::Damage { .. } => {
+                log::debug!("Received Damage event");
+            }
+            zwlr_screencopy_frame_v1::Event::LinuxDmabuf { .. } => {
+                log::debug!("Received LinuxDmaBuf event");
+            }
+            zwlr_screencopy_frame_v1::Event::BufferDone => {
+                log::debug!("Received bufferdone event");
+                frame.buffer_done.store(true, Ordering::SeqCst);
+            }
+            _ => unreachable!(),
+        };
+    }
+}
+
+delegate_noop!(CaptureFrameState: ignore WlShm);
+delegate_noop!(CaptureFrameState: ignore WlShmPool);
+delegate_noop!(CaptureFrameState: ignore WlBuffer);
+delegate_noop!(CaptureFrameState: ignore ZwlrScreencopyManagerV1);
+
+// TODO: Create a xdg-shell surface, check for the enter event, grab the output from it.
+
+pub struct WayshotState {}
+
+impl wayland_client::Dispatch<wl_registry::WlRegistry, GlobalListContents> for WayshotState {
+    fn event(
+        _: &mut WayshotState,
+        _: &wl_registry::WlRegistry,
+        _: wl_registry::Event,
+        _: &GlobalListContents,
+        _: &Connection,
+        _: &QueueHandle<WayshotState>,
+    ) {
+    }
+}

libwayshot/src/error.rs

@@ -0,0 +1,24 @@
+use std::{io, result};
+
+use thiserror::Error;
+use wayland_client::{globals::GlobalError, ConnectError, DispatchError};
+
+pub type Result<T, E = Error> = result::Result<T, E>;
+
+#[derive(Error, Debug)]
+pub enum Error {
+    #[error("no outputs supplied")]
+    NoOutputs,
+    #[error("image buffer is not big enough")]
+    BufferTooSmall,
+    #[error("image color type not supported")]
+    InvalidColor,
+    #[error("IO error: {0}")]
+    Io(#[from] io::Error),
+    #[error("dispatch error: {0}")]
+    Dispatch(#[from] DispatchError),
+    #[error("global error: {0}")]
+    Global(#[from] GlobalError),
+    #[error("connect error: {0}")]
+    Connect(#[from] ConnectError),
+}