Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

tmp #84

Closed
wants to merge 7 commits into from
Closed

tmp #84

Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
Add freeze functionality
  • Loading branch information
AndreasBackx committed Feb 2, 2024
commit 255a685fdc349117e1bf43e770b925c122c01e87
17 changes: 17 additions & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

69 changes: 56 additions & 13 deletions libwayshot/src/dispatch.rs
Original file line number Diff line number Diff line change
@@ -1,20 +1,26 @@
use std::{
collections::HashSet,
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,
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,
};
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,
Expand Down Expand Up @@ -58,16 +64,9 @@ impl Dispatch<WlRegistry, ()> 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.");
Expand Down Expand Up @@ -109,7 +108,11 @@ impl Dispatch<WlOutput, ()> for OutputCaptureState {
} => {
output.transform = transform;
}
_ => (),
wl_output::Event::Scale { factor } => {
output.scale = factor;
}
wl_output::Event::Done => {}
_ => {}
}
}
}
Expand Down Expand Up @@ -227,3 +230,43 @@ impl wayland_client::Dispatch<wl_registry::WlRegistry, GlobalListContents> for W
) {
}
}

pub struct LayerShellState {
pub configured_outputs: HashSet<WlOutput>,
}

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<ZwlrLayerSurfaceV1, WlOutput> for LayerShellState {
// No need to instrument here, span from lib.rs is automatically used.
fn event(
state: &mut Self,
proxy: &ZwlrLayerSurfaceV1,
event: <ZwlrLayerSurfaceV1 as wayland_client::Proxy>::Event,
data: &WlOutput,
_conn: &Connection,
_qhandle: &QueueHandle<Self>,
) {
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")
}
_ => {}
}
}
}
2 changes: 2 additions & 0 deletions libwayshot/src/error.rs
Original file line number Diff line number Diff line change
Expand Up @@ -27,4 +27,6 @@ pub enum Error {
NoSupportedBufferFormat,
#[error("Cannot find required wayland protocol")]
ProtocolNotFound(String),
#[error("error occurred in freeze callback")]
FreezeCallbackError,
}
113 changes: 110 additions & 3 deletions libwayshot/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,13 +12,15 @@ pub mod region;
mod screencopy;

use std::{
collections::HashSet,
fs::File,
os::fd::AsFd,
process::exit,
sync::atomic::{AtomicBool, Ordering},
thread,
};

use dispatch::LayerShellState;
use image::{imageops::replace, DynamicImage};
use memmap2::MmapMut;
use region::{EmbeddedRegion, RegionCapturer};
Expand All @@ -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},
},
Expand All @@ -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::{
Expand Down Expand Up @@ -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<LayerShellState> =
self.conn.new_event_queue::<LayerShellState>();
let qh = event_queue.handle();

let compositor = match self.globals.bind::<WlCompositor, _, _>(&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::<ZwlrLayerShellV1, _, _>(&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,
Expand Down Expand Up @@ -445,13 +535,21 @@ 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)?;

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| {
Expand Down Expand Up @@ -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<dyn Fn() -> Result<LogicalRegion>>,
cursor_overlay: bool,
) -> Result<DynamicImage> {
self.screenshot_region_capturer(RegionCapturer::Freeze(callback), cursor_overlay)
}
/// shot one ouput
pub fn screenshot_single_output(
&self,
Expand Down
7 changes: 4 additions & 3 deletions libwayshot/src/output.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
6 changes: 5 additions & 1 deletion libwayshot/src/region.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand All @@ -12,6 +12,10 @@ pub enum RegionCapturer {
Outputs(Vec<OutputInfo>),
/// 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<dyn Fn() -> Result<LogicalRegion>>),
}

/// `Region` where the coordinate system is the logical coordinate system used
Expand Down
1 change: 1 addition & 0 deletions wayshot/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
4 changes: 2 additions & 2 deletions wayshot/src/clap.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,10 +12,10 @@ pub fn set_flags() -> Command {
.help("Enable debug mode"),
)
.arg(
arg!(-s --slurp <GEOMETRY>)
arg!(-s --slurp <SLURP_ARGS>)
.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 <FILE_PATH>)
Expand Down
Loading
Loading