use anyhow::{Context, Result, bail};
use clap::{Args, Parser, Subcommand, ValueEnum};
use std::io::{Cursor, Write};
use wlr_capture::capture::{self, DEFAULT_BUDGET};
use wlr_capture::{focus, overlay, wl};
#[derive(Parser)]
#[command(
name = "wlr-shot",
version,
about = "Screen capture for wlroots (screenshots, recording, timelapse)"
)]
struct Cli {
#[command(subcommand)]
cmd: Cmd,
}
#[derive(Subcommand)]
enum Cmd {
Screenshot(ShotArgs),
#[cfg(feature = "video")]
Record(RecordArgs),
#[command(hide = true)]
ClipboardServe {
#[arg(long)]
mime: String,
},
}
#[derive(Args)]
struct ShotArgs {
#[arg(short = 's', long, group = "source")]
select: bool,
#[arg(short = 'o', long, value_name = "NAME", group = "source")]
output: Option<String>,
#[arg(short = 'g', long, value_name = "GEOM", group = "source")]
geometry: Option<String>,
#[arg(short = 'w', long, value_name = "ID", group = "source")]
window: Option<String>,
#[arg(long, group = "source")]
pick_window: bool,
#[arg(long, group = "source")]
all: bool,
#[arg(short = 'a', long, group = "source")]
active_window: bool,
#[arg(long, group = "source")]
current_output: bool,
#[arg(short = 't', long, value_enum, default_value_t = Fmt::Png)]
r#type: Fmt,
#[arg(short = 'q', long, default_value_t = 90)]
quality: u8,
#[arg(short = 'c', long)]
clipboard: bool,
#[arg(long, requires = "clipboard")]
clipboard_foreground: bool,
#[arg(long)]
list_outputs: bool,
#[arg(value_name = "FILE", default_value = "-")]
file: String,
}
#[derive(Clone, Copy, ValueEnum)]
enum Fmt {
Png,
Jpeg,
Ppm,
}
pub fn main() {
wlr_capture::i18n::init();
let cli = Cli::parse();
let res = match cli.cmd {
Cmd::Screenshot(args) => screenshot(args),
#[cfg(feature = "video")]
Cmd::Record(args) => record(args),
Cmd::ClipboardServe { mime } => clipboard_serve(&mime),
};
if let Err(e) = res {
eprintln!("wlr-shot: {e:#}");
std::process::exit(1);
}
}
fn screenshot(args: ShotArgs) -> Result<()> {
let mut client = wl::Client::connect().context("Wayland connection")?;
client.refresh().ok();
if args.list_outputs {
for o in client.outputs() {
let (w, h) = o.logical_size();
println!("{}\t{}x{}+{},{}", o.name, w, h, o.logical_x, o.logical_y);
}
return Ok(());
}
let img = if args.select {
let conn = wlr_capture::Connection::connect_to_env()?;
let caps = capture::capture_all(&mut client, DEFAULT_BUDGET)?;
match overlay::select_region_on(&conn, &caps)? {
Some(region) => capture::composite(&caps, region)?,
None => std::process::exit(1), }
} else if let Some(geo) = &args.geometry {
capture::capture_region(&mut client, capture::parse_geometry(geo)?, DEFAULT_BUDGET)?
} else if args.all {
let region = capture::whole_layout(&client)?;
capture::capture_region(&mut client, region, DEFAULT_BUDGET)?
} else if args.active_window {
capture::capture_region(&mut client, active_window_rect()?, DEFAULT_BUDGET)?
} else if args.current_output {
capture::capture_output(&mut client, Some(&focused_output()?), DEFAULT_BUDGET)?
} else if args.pick_window {
capture::capture_window(&mut client, &pick_window()?, DEFAULT_BUDGET)?
} else if let Some(id) = &args.window {
capture::capture_window(&mut client, id, DEFAULT_BUDGET)?
} else {
capture::capture_output(&mut client, args.output.as_deref(), DEFAULT_BUDGET)?
};
let bytes = encode(&img, args.r#type, args.quality).context("encoding image")?;
if args.clipboard {
clipboard_copy(mime_for(args.r#type), bytes, args.clipboard_foreground)
.context("copying to clipboard")?;
} else {
write_out(&args.file, &bytes).context("writing output")?;
}
Ok(())
}
fn mime_for(fmt: Fmt) -> &'static str {
match fmt {
Fmt::Png => "image/png",
Fmt::Jpeg => "image/jpeg",
Fmt::Ppm => "image/x-portable-pixmap",
}
}
fn clipboard_copy(mime: &str, bytes: Vec<u8>, foreground: bool) -> Result<()> {
if foreground {
return wlr_capture::clipboard::serve(mime, bytes);
}
wlr_capture::clipboard::spawn_detached(&bytes, &["clipboard-serve", "--mime", mime])
}
fn clipboard_serve(mime: &str) -> Result<()> {
use std::io::Read;
let mut data = Vec::new();
std::io::stdin()
.read_to_end(&mut data)
.context("reading clipboard data")?;
wlr_capture::clipboard::serve(mime, data)
}
fn active_window_rect() -> Result<wl::Region> {
let backend = focus::detect()
.context("focus info unavailable (unsupported compositor); try --pick-window")?;
backend
.active_window_rect()
.with_context(|| format!("no active window detected (via {})", backend.name()))
}
fn focused_output() -> Result<String> {
let backend = focus::detect()
.context("focus info unavailable (unsupported compositor); specify -o NAME")?;
backend
.focused_output()
.with_context(|| format!("no focused output detected (via {})", backend.name()))
}
fn pick_window() -> Result<String> {
let sibling = std::env::current_exe()
.ok()
.and_then(|p| p.parent().map(|d| d.join("wlr-chooser")))
.filter(|p| p.exists());
let mut cmd = match sibling {
Some(p) => std::process::Command::new(p),
None => std::process::Command::new("wlr-chooser"),
};
let out = cmd
.arg("--windows")
.output()
.context("lancement de wlr-chooser")?;
String::from_utf8_lossy(&out.stdout)
.lines()
.find_map(|l| l.strip_prefix("Window: "))
.map(|s| s.to_string())
.context("no window picked")
}
fn encode(img: &wl::CapturedImage, fmt: Fmt, quality: u8) -> Result<Vec<u8>> {
use image::{DynamicImage, ImageFormat, RgbaImage, codecs::jpeg::JpegEncoder};
if img.width == 0 || img.height == 0 {
bail!("empty image (region off-screen?)");
}
let rgba = RgbaImage::from_raw(img.width, img.height, img.rgba.clone())
.context("inconsistent dimensions/buffer")?;
let dynimg = DynamicImage::ImageRgba8(rgba);
let mut out = Vec::new();
let mut cur = Cursor::new(&mut out);
match fmt {
Fmt::Png => dynimg.write_to(&mut cur, ImageFormat::Png)?,
Fmt::Ppm => {
DynamicImage::ImageRgb8(dynimg.to_rgb8()).write_to(&mut cur, ImageFormat::Pnm)?
}
Fmt::Jpeg => {
JpegEncoder::new_with_quality(&mut cur, quality.clamp(1, 100))
.encode_image(&dynimg.to_rgb8())?;
}
}
Ok(out)
}
fn write_out(file: &str, bytes: &[u8]) -> Result<()> {
if file == "-" {
std::io::stdout().write_all(bytes)?;
} else {
std::fs::write(file, bytes)?;
}
Ok(())
}
#[cfg(feature = "video")]
mod record_impl {
use super::{active_window_rect, focused_output, pick_window};
use anyhow::{Context, Result, bail};
use clap::{Args, ValueEnum};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::{Duration, Instant};
use wlr_capture::capture;
use wlr_capture::gl::GpuReadback;
use wlr_capture::sink::FrameSink;
use wlr_capture::stream;
use wlr_capture::video::{self, VideoEncoder};
use wlr_capture::wl::{self, CapturedImage, Output, Region};
const ROUND: Duration = Duration::from_millis(200);
const APPEAR_GRACE: Duration = Duration::from_secs(5);
const GIF_MAX_DIM: u32 = 800;
const WEBP_MAX_DIM: u32 = 1280;
fn downscaled(img: &CapturedImage, max_dim: u32) -> Result<image::RgbaImage> {
let buf = image::RgbaImage::from_raw(img.width, img.height, img.rgba.clone())
.ok_or_else(|| anyhow::anyhow!("image dimensions don't match the buffer"))?;
let long = img.width.max(img.height);
if long <= max_dim {
return Ok(buf);
}
let scale = max_dim as f32 / long as f32;
let nw = ((img.width as f32 * scale).round() as u32).max(1);
let nh = ((img.height as f32 * scale).round() as u32).max(1);
Ok(image::imageops::resize(
&buf,
nw,
nh,
image::imageops::FilterType::Triangle,
))
}
struct GifSink {
enc: image::codecs::gif::GifEncoder<std::io::BufWriter<std::fs::File>>,
delay: image::Delay,
}
impl GifSink {
fn new(path: &str, fps: u32) -> Result<Self> {
use image::codecs::gif::{GifEncoder, Repeat};
let file = std::fs::File::create(path).context("creating the GIF file")?;
let mut enc = GifEncoder::new_with_speed(std::io::BufWriter::new(file), 30);
enc.set_repeat(Repeat::Infinite).context("GIF loop flag")?;
let ms = (1000.0 / fps as f64).round().max(20.0) as u32;
Ok(Self {
enc,
delay: image::Delay::from_numer_denom_ms(ms, 1),
})
}
}
impl FrameSink for GifSink {
fn push(&mut self, img: &CapturedImage, _ts: Duration) -> Result<()> {
let buf = downscaled(img, GIF_MAX_DIM)?;
self.enc
.encode_frame(image::Frame::from_parts(buf, 0, 0, self.delay))
.context("encoding a GIF frame")?;
Ok(())
}
}
struct WebpSink {
path: String,
fps: u32,
dims: Option<(u32, u32)>,
frames: Vec<Vec<u8>>,
}
impl WebpSink {
fn new(path: &str, fps: u32) -> Self {
Self {
path: path.to_string(),
fps,
dims: None,
frames: Vec::new(),
}
}
}
impl FrameSink for WebpSink {
fn push(&mut self, img: &CapturedImage, _ts: Duration) -> Result<()> {
let buf = downscaled(img, WEBP_MAX_DIM)?;
let dims = *self.dims.get_or_insert((buf.width(), buf.height()));
if (buf.width(), buf.height()) == dims {
self.frames.push(buf.into_raw());
}
Ok(())
}
fn finish(&mut self) -> Result<()> {
let Some((w, h)) = self.dims else {
return Ok(());
};
let mut config = webp::WebPConfig::new().map_err(|_| anyhow::anyhow!("WebP config"))?;
config.quality = 75.0;
let mut enc = webp::AnimEncoder::new(w, h, &config);
let step = (1000.0 / self.fps as f64).round() as i32;
for (i, rgba) in self.frames.iter().enumerate() {
enc.add_frame(webp::AnimFrame::from_rgba(rgba, w, h, i as i32 * step));
}
let data = enc
.try_encode()
.map_err(|e| anyhow::anyhow!("WebP encode: {e:?}"))?;
std::fs::write(&self.path, &*data).context("writing the WebP file")?;
Ok(())
}
}
fn is_video_ext(file: &str) -> bool {
let ext = std::path::Path::new(file)
.extension()
.and_then(|e| e.to_str())
.unwrap_or("")
.to_ascii_lowercase();
!matches!(ext.as_str(), "gif" | "webp")
}
fn make_sink(
args: &RecordArgs,
mode: video::Mode,
audio: bool,
) -> Result<(Box<dyn FrameSink>, String)> {
let fps = args.fps.max(1);
if !is_video_ext(&args.file) {
let ext = std::path::Path::new(&args.file)
.extension()
.and_then(|e| e.to_str())
.unwrap_or("")
.to_ascii_lowercase();
return match ext.as_str() {
"gif" => Ok((Box::new(GifSink::new(&args.file, fps)?), "GIF".into())),
_ => Ok((Box::new(WebpSink::new(&args.file, fps)), "WebP".into())),
};
}
let enc = VideoEncoder::new(
&args.file,
video::Options {
backend: args.encoder.into(),
fps,
mode,
device: Some(args.device.clone().into()),
audio,
},
)?;
let label = format!("{:?}", enc.resolved_backend()?);
Ok((Box::new(enc), label))
}
#[cfg(feature = "audio")]
fn audio_wanted(args: &RecordArgs) -> bool {
!args.no_audio && is_video_ext(&args.file)
}
#[derive(Args)]
pub struct RecordArgs {
#[arg(short = 's', long, group = "source")]
select: bool,
#[arg(short = 'o', long, value_name = "NAME", group = "source")]
output: Option<String>,
#[arg(short = 'g', long, value_name = "GEOM", group = "source")]
geometry: Option<String>,
#[arg(short = 'w', long, value_name = "ID", group = "source")]
window: Option<String>,
#[arg(long, group = "source")]
pick_window: bool,
#[arg(short = 'a', long, group = "source")]
active_window: bool,
#[arg(long, group = "source")]
current_output: bool,
#[arg(long, value_enum, default_value_t = Enc::Auto)]
encoder: Enc,
#[arg(long, value_name = "PATH", default_value = "/dev/dri/renderD128")]
device: String,
#[arg(long, default_value_t = 30)]
fps: u32,
#[arg(long, value_name = "INTERVAL")]
timelapse: Option<String>,
#[arg(short = 'd', long, value_name = "SECS")]
duration: Option<f64>,
#[cfg(feature = "audio")]
#[arg(long)]
no_audio: bool,
#[cfg(feature = "audio")]
#[arg(long, value_name = "NODE")]
audio_source: Option<String>,
#[arg(value_name = "FILE")]
file: String,
}
#[derive(Clone, Copy, ValueEnum)]
pub enum Enc {
Auto,
Nvenc,
Vaapi,
Software,
}
impl From<Enc> for video::Backend {
fn from(e: Enc) -> Self {
match e {
Enc::Auto => video::Backend::Auto,
Enc::Nvenc => video::Backend::Nvenc,
Enc::Vaapi => video::Backend::Vaapi,
Enc::Software => video::Backend::Software,
}
}
}
enum Target {
Output {
output: Output,
crop: Option<Region>,
},
Window(String),
}
impl Target {
fn label(&self) -> String {
match self {
Target::Output { output, crop: None } => format!("output {}", output.name),
Target::Output {
output,
crop: Some(c),
} => {
format!("region {}x{} on {}", c.w, c.h, output.name)
}
Target::Window(id) => format!("window {id}"),
}
}
}
fn parse_interval(s: &str) -> Result<Duration> {
let s = s.trim();
let err = || anyhow::anyhow!("invalid interval '{s}' (try e.g. 2s, 500ms, 1m)");
let (num, mult) = if let Some(n) = s.strip_suffix("ms") {
(n, 0.001)
} else if let Some(n) = s.strip_suffix('s') {
(n, 1.0)
} else if let Some(n) = s.strip_suffix('m') {
(n, 60.0)
} else {
(s, 1.0) };
let secs: f64 = num.trim().parse().map_err(|_| err())?;
if !(secs.is_finite() && secs > 0.0) {
return Err(err());
}
Ok(Duration::from_secs_f64(secs * mult))
}
fn region_target(client: &wl::Client, region: Region) -> Result<Target> {
if region.is_empty() {
bail!("empty region");
}
let corner = Region {
x: region.x,
y: region.y,
w: 1,
h: 1,
};
let output = client
.outputs()
.iter()
.find(|o| o.logical_rect().intersect(&corner).is_some())
.cloned()
.context("the region's top-left corner is on no output")?;
let clipped = region
.intersect(&output.logical_rect())
.context("region does not overlap its output")?;
if clipped != region {
eprintln!(
"wlr-shot: region spans multiple outputs; recording the {}x{} part on {}",
clipped.w, clipped.h, output.name
);
}
let crop = capture::logical_to_physical(&output, clipped);
Ok(Target::Output {
output,
crop: Some(crop),
})
}
fn resolve_target(client: &mut wl::Client, args: &RecordArgs) -> Result<Target> {
if args.select {
let conn = wlr_capture::Connection::connect_to_env()?;
let caps = capture::capture_all(client, capture::DEFAULT_BUDGET)?;
match wlr_capture::overlay::select_region_on(&conn, &caps)? {
Some(region) => region_target(client, region),
None => std::process::exit(1), }
} else if let Some(geo) = &args.geometry {
region_target(client, capture::parse_geometry(geo)?)
} else if args.active_window {
region_target(client, active_window_rect()?)
} else if let Some(id) = &args.window {
Ok(Target::Window(id.clone()))
} else if args.pick_window {
Ok(Target::Window(pick_window()?))
} else {
let name = if args.current_output {
Some(focused_output()?)
} else {
args.output.clone()
};
Ok(Target::Output {
output: resolve_output(client, name.as_deref())?,
crop: None,
})
}
}
fn resolve_output(client: &wl::Client, name: Option<&str>) -> Result<Output> {
let outputs = client.outputs();
match name {
Some(n) => outputs
.iter()
.find(|o| o.name == n)
.cloned()
.with_context(|| format!("output '{n}' not found")),
None => match outputs {
[single] => Ok(single.clone()),
[] => bail!("no outputs available"),
many => {
let names: Vec<&str> = many.iter().map(|o| o.name.as_str()).collect();
bail!(
"multiple outputs; specify -o NAME among: {}",
names.join(", ")
)
}
},
}
}
pub fn record(args: RecordArgs) -> Result<()> {
let mut client = wl::Client::connect().context("Wayland connection")?;
client.refresh().ok();
let target = resolve_target(&mut client, &args)?;
let mode = match &args.timelapse {
Some(_) => video::Mode::Timelapse,
None => video::Mode::Record,
};
let interval = args.timelapse.as_deref().map(parse_interval).transpose()?;
#[cfg(feature = "audio")]
let audio = if audio_wanted(&args) {
match wlr_capture::audio::AudioCapture::start(args.audio_source.clone()) {
Ok(c) => Some(c),
Err(e) => {
eprintln!("wlr-shot: recording without audio ({e})");
None
}
}
} else {
None
};
#[cfg(feature = "audio")]
let have_audio = audio.is_some();
#[cfg(not(feature = "audio"))]
let have_audio = false;
let (mut sink, fmt) = make_sink(&args, mode, have_audio)?;
eprintln!(
"wlr-shot: recording {} to {} ({fmt}{}). Press Ctrl-C to stop.",
target.label(),
args.file,
if have_audio { " + audio" } else { "" },
);
let stop = Arc::new(AtomicBool::new(false));
let s = stop.clone();
ctrlc::set_handler(move || s.store(true, Ordering::SeqCst))
.context("installing Ctrl-C handler")?;
let start = Instant::now();
let deadline = args.duration.map(|d| start + Duration::from_secs_f64(d));
let crop = match &target {
Target::Output { crop, .. } => *crop,
Target::Window(_) => None,
};
let stream_source = match &target {
Target::Output { output, .. } => stream::Source::Output(output.name.clone()),
Target::Window(id) => stream::Source::Toplevel(id.clone()),
};
let frame_interval =
interval.unwrap_or_else(|| Duration::from_secs_f64(1.0 / args.fps.max(1) as f64));
let mut rb: Option<GpuReadback> = None;
let mut stream = stream::Stream::new(stream_source, APPEAR_GRACE);
let mut frames = 0u64;
let mut last_img: Option<CapturedImage> = None; let mut next_tick: Option<Duration> = None; let mut last_log = start;
loop {
if stop.load(Ordering::SeqCst) {
break;
}
if let Some(dl) = deadline
&& Instant::now() >= dl
{
break;
}
let budget = match next_tick {
Some(nt) => nt
.saturating_sub(start.elapsed())
.clamp(Duration::from_millis(1), ROUND),
None => ROUND,
};
let step = stream.step(&mut client, budget);
for frame in step.frames {
let mut img = stream::decode_frame(&mut rb, frame)?;
if let Some(c) = crop {
img = img.crop(c);
}
last_img = Some(img);
next_tick.get_or_insert_with(|| start.elapsed());
}
if let Some(end) = step.end {
match end {
stream::End::NeverAppeared => {
bail!("source did not appear within {}s", APPEAR_GRACE.as_secs())
}
stream::End::SourceGone => break, }
}
#[cfg(feature = "audio")]
if let Some(cap) = &audio {
let pcm = cap.drain();
if !pcm.is_empty() {
sink.push_audio(&pcm);
}
}
if let (Some(img), Some(mut nt)) = (last_img.as_ref(), next_tick) {
while start.elapsed() >= nt {
sink.push(img, nt)?; frames += 1;
nt += frame_interval;
}
next_tick = Some(nt);
}
if last_log.elapsed() >= Duration::from_secs(1) {
eprint!(
"\rwlr-shot: {frames} frames, {:.0}s ",
start.elapsed().as_secs_f64()
);
std::io::Write::flush(&mut std::io::stderr()).ok();
last_log = Instant::now();
}
}
#[cfg(feature = "audio")]
if let Some(cap) = &audio {
let pcm = cap.drain();
if !pcm.is_empty() {
sink.push_audio(&pcm);
}
}
#[cfg(feature = "audio")]
drop(audio);
sink.finish().context("finalising the output file")?;
drop(sink); eprintln!(
"\rwlr-shot: saved {} ({frames} frames, {:.1}s) ",
args.file,
start.elapsed().as_secs_f64()
);
Ok(())
}
}
#[cfg(feature = "video")]
use record_impl::{RecordArgs, record};