Capture terminal color palette

This uses OSC sequence to query the colors from the terminal, and saves
them as a theme in asciicast v2 header
(https://docs.asciinema.org/manual/asciicast/v2/#theme).
This commit is contained in:
Marcin Kulik
2024-02-12 15:49:35 +01:00
parent 1d0b7da09c
commit bf52290514
10 changed files with 381 additions and 37 deletions

1
Cargo.lock generated
View File

@@ -93,6 +93,7 @@ dependencies = [
"mime_guess",
"nix",
"reqwest",
"rgb",
"rust-embed",
"rustyline",
"scraper",

View File

@@ -35,3 +35,4 @@ mime_guess = "2.0.4"
tower-http = { version = "0.5.1", features = ["trace"] }
tracing = "0.1.40"
tracing-subscriber = { version = "0.3.18", features = ["env-filter"] }
rgb = "0.8.37"

View File

@@ -1,6 +1,7 @@
mod util;
mod v1;
mod v2;
use crate::tty;
use anyhow::{anyhow, Result};
use std::collections::HashMap;
use std::fs;
@@ -22,6 +23,7 @@ pub struct Header {
pub command: Option<String>,
pub title: Option<String>,
pub env: Option<HashMap<String, String>>,
pub theme: Option<tty::Theme>,
}
pub struct Event {
@@ -137,7 +139,9 @@ pub fn accelerate(
#[cfg(test)]
mod tests {
use super::{Asciicast, Event, EventData, Header, Writer};
use crate::tty;
use anyhow::Result;
use rgb::RGB8;
use std::collections::HashMap;
use std::io;
@@ -150,6 +154,7 @@ mod tests {
assert_eq!(header.version, 1);
assert_eq!((header.cols, header.rows), (100, 50));
assert!(header.theme.is_none());
assert_eq!(events[0].time, 1230000);
assert!(matches!(events[0].data, EventData::Output(ref s) if s == "hello"));
@@ -174,20 +179,45 @@ mod tests {
}
#[test]
fn open_v2() {
let Asciicast { header, events } = super::open_from_path("tests/casts/demo.cast").unwrap();
let events = events.take(7).collect::<Result<Vec<Event>>>().unwrap();
fn open_v2_minimal() {
let Asciicast { header, events } =
super::open_from_path("tests/casts/minimal.cast").unwrap();
let events = events.collect::<Result<Vec<Event>>>().unwrap();
assert_eq!((header.cols, header.rows), (75, 18));
assert_eq!((header.cols, header.rows), (100, 50));
assert!(header.theme.is_none());
assert_eq!(events[1].time, 100989);
assert!(matches!(events[1].data, EventData::Output(ref s) if s == "\u{1b}[?2004h"));
assert_eq!(events[0].time, 1230000);
assert!(matches!(events[0].data, EventData::Output(ref s) if s == "hello"));
}
assert_eq!(events[5].time, 1511526);
assert!(matches!(events[5].data, EventData::Input(ref s) if s == "v"));
#[test]
fn open_v2_full() {
let Asciicast { header, events } = super::open_from_path("tests/casts/full.cast").unwrap();
let events = events.take(5).collect::<Result<Vec<Event>>>().unwrap();
let theme = header.theme.unwrap();
assert_eq!(events[6].time, 1511937);
assert!(matches!(events[6].data, EventData::Output(ref s) if s == "v"));
assert_eq!((header.cols, header.rows), (100, 50));
assert_eq!(theme.fg, RGB8::new(0, 0, 0));
assert_eq!(theme.bg, RGB8::new(0xff, 0xff, 0xff));
assert_eq!(theme.palette[0], RGB8::new(0x24, 0x1f, 0x31));
assert_eq!(events[0].time, 1);
assert!(matches!(events[0].data, EventData::Output(ref s) if s == "ż"));
assert_eq!(events[1].time, 1_000_000);
assert!(matches!(events[1].data, EventData::Output(ref s) if s == "ółć"));
assert_eq!(events[2].time, 2_300_000);
assert!(matches!(events[2].data, EventData::Input(ref s) if s == "\n"));
assert_eq!(events[3].time, 5_600_001);
assert!(
matches!(events[3].data, EventData::Resize(ref cols, ref rows) if *cols == 80 && *rows == 40)
);
assert_eq!(events[4].time, 10_500_000);
assert!(matches!(events[4].data, EventData::Output(ref s) if s == "\r\n"));
}
#[test]
@@ -206,6 +236,7 @@ mod tests {
command: None,
title: None,
env: Default::default(),
theme: None,
};
fw.write_header(&header).unwrap();
@@ -262,6 +293,29 @@ mod tests {
env.insert("SHELL".to_owned(), "/usr/bin/fish".to_owned());
env.insert("TERM".to_owned(), "xterm256-color".to_owned());
let theme = tty::Theme {
fg: RGB8::new(0, 1, 2),
bg: RGB8::new(0, 100, 200),
palette: vec![
RGB8::new(0, 0, 0),
RGB8::new(10, 11, 12),
RGB8::new(20, 21, 22),
RGB8::new(30, 31, 32),
RGB8::new(40, 41, 42),
RGB8::new(50, 51, 52),
RGB8::new(60, 61, 62),
RGB8::new(70, 71, 72),
RGB8::new(80, 81, 82),
RGB8::new(90, 91, 92),
RGB8::new(100, 101, 102),
RGB8::new(110, 111, 112),
RGB8::new(120, 121, 122),
RGB8::new(130, 131, 132),
RGB8::new(140, 141, 142),
RGB8::new(150, 151, 152),
],
};
let header = Header {
version: 2,
cols: 80,
@@ -271,6 +325,7 @@ mod tests {
command: Some("/bin/bash".to_owned()),
title: Some("Demo".to_owned()),
env: Some(env),
theme: Some(theme),
};
fw.write_header(&header).unwrap();
@@ -288,6 +343,9 @@ mod tests {
assert_eq!(lines[0]["env"].as_object().unwrap().len(), 2);
assert_eq!(lines[0]["env"]["SHELL"], "/usr/bin/fish");
assert_eq!(lines[0]["env"]["TERM"], "xterm256-color");
assert_eq!(lines[0]["theme"]["fg"], "#000102");
assert_eq!(lines[0]["theme"]["bg"], "#0064c8");
assert_eq!(lines[0]["theme"]["palette"], "#000000:#0a0b0c:#141516:#1e1f20:#28292a:#323334:#3c3d3e:#464748:#505152:#5a5b5c:#646566:#6e6f70:#78797a:#828384:#8c8d8e:#969798");
}
fn parse(json: Vec<u8>) -> Vec<serde_json::Value> {

View File

@@ -38,6 +38,7 @@ pub fn load(json: String) -> Result<Asciicast<'static>> {
command: asciicast.command.clone(),
title: asciicast.title.clone(),
env: asciicast.env.clone(),
theme: None,
};
let events = Box::new(

View File

@@ -1,6 +1,7 @@
use super::{util, Asciicast, Event, EventData, Header};
use crate::tty;
use anyhow::{anyhow, bail, Result};
use serde::{Deserialize, Deserializer};
use serde::{Deserialize, Deserializer, Serialize};
use std::collections::HashMap;
use std::io::{self, Write};
@@ -14,8 +15,25 @@ struct V2Header {
command: Option<String>,
title: Option<String>,
env: Option<HashMap<String, String>>,
theme: Option<V2Theme>,
}
#[derive(Deserialize, Serialize, Clone)]
struct V2Theme {
#[serde(deserialize_with = "deserialize_color")]
fg: RGB8,
#[serde(deserialize_with = "deserialize_color")]
bg: RGB8,
#[serde(deserialize_with = "deserialize_palette")]
palette: V2Palette,
}
#[derive(Clone)]
struct RGB8(rgb::RGB8);
#[derive(Clone)]
struct V2Palette(Vec<RGB8>);
#[derive(Debug, Deserialize)]
struct V2Event {
#[serde(deserialize_with = "util::deserialize_time")]
@@ -60,6 +78,7 @@ impl Parser {
command: self.0.command.clone(),
title: self.0.title.clone(),
env: self.0.env.clone(),
theme: self.0.theme.as_ref().map(|t| t.into()),
};
let events = Box::new(lines.filter_map(parse_line));
@@ -215,6 +234,10 @@ impl serde::Serialize for V2Header {
len += 1;
}
if self.theme.is_some() {
len += 1;
}
let mut map = serializer.serialize_map(Some(len))?;
map.serialize_entry("version", &2)?;
map.serialize_entry("width", &self.width)?;
@@ -242,10 +265,82 @@ impl serde::Serialize for V2Header {
}
}
if let Some(theme) = &self.theme {
map.serialize_entry("theme", &theme)?;
}
map.end()
}
}
fn deserialize_color<'de, D>(deserializer: D) -> Result<RGB8, D::Error>
where
D: Deserializer<'de>,
{
let value: &str = Deserialize::deserialize(deserializer)?;
parse_hex_color(value).ok_or(serde::de::Error::custom("invalid hex triplet"))
}
fn parse_hex_color(rgb: &str) -> Option<RGB8> {
if rgb.len() != 7 {
return None;
}
let r = u8::from_str_radix(&rgb[1..3], 16).ok()?;
let g = u8::from_str_radix(&rgb[3..5], 16).ok()?;
let b = u8::from_str_radix(&rgb[5..7], 16).ok()?;
Some(RGB8(rgb::RGB8::new(r, g, b)))
}
fn deserialize_palette<'de, D>(deserializer: D) -> Result<V2Palette, D::Error>
where
D: Deserializer<'de>,
{
let value: &str = Deserialize::deserialize(deserializer)?;
let mut colors: Vec<RGB8> = value.split(':').filter_map(parse_hex_color).collect();
let len = colors.len();
if len == 8 {
colors.extend_from_within(..);
} else if len != 16 {
return Err(serde::de::Error::custom("expected 8 or 16 hex triplets"));
}
Ok(V2Palette(colors))
}
impl serde::Serialize for RGB8 {
fn serialize<S>(&self, serializer: S) -> std::prelude::v1::Result<S::Ok, S::Error>
where
S: serde::Serializer,
{
serializer.serialize_str(&self.to_string())
}
}
impl ToString for RGB8 {
fn to_string(&self) -> String {
format!("#{:0>2x}{:0>2x}{:0>2x}", self.0.r, self.0.g, self.0.b)
}
}
impl serde::Serialize for V2Palette {
fn serialize<S>(&self, serializer: S) -> std::prelude::v1::Result<S::Ok, S::Error>
where
S: serde::Serializer,
{
let palette = self
.0
.iter()
.map(|c| c.to_string())
.collect::<Vec<_>>()
.join(":");
serializer.serialize_str(&palette)
}
}
impl From<&Header> for V2Header {
fn from(header: &Header) -> Self {
V2Header {
@@ -257,6 +352,31 @@ impl From<&Header> for V2Header {
command: header.command.clone(),
title: header.title.clone(),
env: header.env.clone(),
theme: header.theme.as_ref().map(|t| t.into()),
}
}
}
impl From<&tty::Theme> for V2Theme {
fn from(theme: &tty::Theme) -> Self {
let palette = theme.palette.iter().copied().map(RGB8).collect();
V2Theme {
fg: RGB8(theme.fg),
bg: RGB8(theme.bg),
palette: V2Palette(palette),
}
}
}
impl From<&V2Theme> for tty::Theme {
fn from(theme: &V2Theme) -> Self {
let palette = theme.palette.0.iter().map(|c| c.0).collect();
tty::Theme {
fg: theme.fg.0,
bg: theme.bg.0,
palette,
}
}
}

View File

@@ -63,7 +63,7 @@ pub struct Cli {
rows: Option<u16>,
}
#[derive(Clone, Copy, Debug, ValueEnum)]
#[derive(Clone, Copy, Debug, PartialEq, ValueEnum)]
enum Format {
Asciicast,
Raw,
@@ -74,10 +74,11 @@ impl Cli {
pub fn run(self, config: &Config) -> Result<()> {
locale::check_utf8_locale()?;
let format = self.get_format();
let (append, overwrite) = self.get_mode()?;
let file = self.open_file(append, overwrite)?;
let time_offset = self.get_time_offset(append, format)?;
let command = self.get_command(config);
let output = self.get_output(file, append, config)?;
let keys = get_key_bindings(config)?;
let notifier = super::get_notifier(config);
let record_input = self.input || config.cmd_rec_input();
@@ -98,6 +99,8 @@ impl Cli {
Box::new(tty::NullTty::open()?)
};
let theme = tty.get_theme();
let output = self.get_output(file, format, append, time_offset, theme, config);
let mut recorder = recorder::Recorder::new(output, record_input, keys, notifier);
pty::exec(
@@ -149,13 +152,8 @@ impl Cli {
Ok(file)
}
fn get_output(
&self,
file: fs::File,
append: bool,
config: &Config,
) -> Result<Box<dyn recorder::Output + Send>> {
let format = self.format.unwrap_or_else(|| {
fn get_format(&self) -> Format {
self.format.unwrap_or_else(|| {
if self.raw {
Format::Raw
} else if self.filename.to_lowercase().ends_with(".txt") {
@@ -163,28 +161,40 @@ impl Cli {
} else {
Format::Asciicast
}
});
})
}
fn get_time_offset(&self, append: bool, format: Format) -> Result<u64> {
if append && format == Format::Asciicast {
asciicast::get_duration(&self.filename)
} else {
Ok(0)
}
}
fn get_output(
&self,
file: fs::File,
format: Format,
append: bool,
time_offset: u64,
theme: Option<tty::Theme>,
config: &Config,
) -> Box<dyn recorder::Output + Send> {
match format {
Format::Asciicast => {
let time_offset = if append {
asciicast::get_duration(&self.filename)?
} else {
0
};
let metadata = self.build_asciicast_metadata(theme, config);
let metadata = self.build_asciicast_metadata(config);
Ok(Box::new(encoder::AsciicastEncoder::new(
Box::new(encoder::AsciicastEncoder::new(
file,
append,
time_offset,
metadata,
)))
))
}
Format::Raw => Ok(Box::new(encoder::RawEncoder::new(file, append))),
Format::Txt => Ok(Box::new(encoder::TextEncoder::new(file))),
Format::Raw => Box::new(encoder::RawEncoder::new(file, append)),
Format::Txt => Box::new(encoder::TextEncoder::new(file)),
}
}
@@ -192,7 +202,11 @@ impl Cli {
self.command.as_ref().cloned().or(config.cmd_rec_command())
}
fn build_asciicast_metadata(&self, config: &Config) -> encoder::Metadata {
fn build_asciicast_metadata(
&self,
theme: Option<tty::Theme>,
config: &Config,
) -> encoder::Metadata {
let idle_time_limit = self.idle_time_limit.or(config.cmd_rec_idle_time_limit());
let command = self.get_command(config);
@@ -208,6 +222,7 @@ impl Cli {
command,
title: self.title.clone(),
env: Some(capture_env(&env)),
theme,
}
}
}

View File

@@ -14,6 +14,7 @@ pub struct Metadata {
pub command: Option<String>,
pub title: Option<String>,
pub env: Option<HashMap<String, String>>,
pub theme: Option<tty::Theme>,
}
impl<W> AsciicastEncoder<W>
@@ -38,6 +39,7 @@ where
command: self.metadata.command.clone(),
title: self.metadata.title.clone(),
env: self.metadata.env.clone(),
theme: self.metadata.theme.clone(),
}
}
}
@@ -68,6 +70,7 @@ impl From<&Header> for Metadata {
command: header.command.as_ref().cloned(),
title: header.title.as_ref().cloned(),
env: header.env.as_ref().cloned(),
theme: header.theme.as_ref().cloned(),
}
}
}

View File

@@ -1,9 +1,17 @@
use anyhow::Result;
use nix::{libc, pty, unistd};
use std::{
fs, io,
os::fd::{AsFd, AsRawFd, BorrowedFd, FromRawFd, OwnedFd},
use nix::{
errno::Errno,
libc, pty,
sys::{
select::{select, FdSet},
time::TimeVal,
},
unistd,
};
use rgb::RGB8;
use std::fs;
use std::io;
use std::os::fd::{AsFd, AsRawFd, BorrowedFd, FromRawFd, OwnedFd};
use termion::raw::{IntoRawMode, RawTerminal};
#[derive(Clone, Copy, Debug, PartialEq)]
@@ -23,6 +31,14 @@ impl From<TtySize> for (u16, u16) {
pub trait Tty: io::Write + io::Read + AsFd {
fn get_size(&self) -> pty::Winsize;
fn get_theme(&self) -> Option<Theme>;
}
#[derive(Clone)]
pub struct Theme {
pub fg: RGB8,
pub bg: RGB8,
pub palette: Vec<RGB8>,
}
pub struct DevTty {
@@ -43,6 +59,25 @@ impl DevTty {
}
}
fn parse_color(rgb: &str) -> Option<RGB8> {
let mut components = rgb.split('/');
let r_hex = components.next()?;
let g_hex = components.next()?;
let b_hex = components.next()?;
if r_hex.len() < 2 || g_hex.len() < 2 || b_hex.len() < 2 {
return None;
}
let r = u8::from_str_radix(&r_hex[..2], 16).ok()?;
let g = u8::from_str_radix(&g_hex[..2], 16).ok()?;
let b = u8::from_str_radix(&b_hex[..2], 16).ok()?;
Some(RGB8::new(r, g, b))
}
static COLORS_QUERY: &[u8; 148] = b"\x1b]10;?\x07\x1b]11;?\x07\x1b]4;0;?\x07\x1b]4;1;?\x07\x1b]4;2;?\x07\x1b]4;3;?\x07\x1b]4;4;?\x07\x1b]4;5;?\x07\x1b]4;6;?\x07\x1b]4;7;?\x07\x1b]4;8;?\x07\x1b]4;9;?\x07\x1b]4;10;?\x07\x1b]4;11;?\x07\x1b]4;12;?\x07\x1b]4;13;?\x07\x1b]4;14;?\x07\x1b]4;15;?\x07";
impl Tty for DevTty {
fn get_size(&self) -> pty::Winsize {
let mut winsize = pty::Winsize {
@@ -56,6 +91,74 @@ impl Tty for DevTty {
winsize
}
fn get_theme(&self) -> Option<Theme> {
let mut query = &COLORS_QUERY[..];
let mut response = Vec::new();
let mut buf = [0u8; 1024];
let mut color_count = 0;
let fd = self.as_fd().as_raw_fd();
loop {
let mut timeout = TimeVal::new(0, 50_000);
let mut rfds = FdSet::new();
let mut wfds = FdSet::new();
rfds.insert(self);
if !query.is_empty() {
wfds.insert(self);
}
match select(None, &mut rfds, &mut wfds, None, &mut timeout) {
Ok(0) => return None,
Ok(_) => {
if rfds.contains(self) {
let n = unistd::read(fd, &mut buf).ok()?;
response.extend_from_slice(&buf[..n]);
color_count += &buf[..n]
.iter()
.filter(|b| *b == &0x07 || *b == &b'\\')
.count();
if color_count == 18 {
break;
}
}
if wfds.contains(self) {
let n = unistd::write(fd, query).ok()?;
query = &query[n..];
}
}
Err(e) => {
if e == Errno::EINTR {
continue;
} else {
return None;
}
}
}
}
let response = String::from_utf8_lossy(response.as_slice());
let mut colors = response.match_indices("rgb:");
let (idx, _) = colors.next()?;
let fg = parse_color(&response[idx + 4..])?;
let (idx, _) = colors.next()?;
let bg = parse_color(&response[idx + 4..])?;
let mut palette = Vec::new();
for _ in 0..16 {
let (idx, _) = colors.next()?;
let color = parse_color(&response[idx + 4..])?;
palette.push(color);
}
Some(Theme { fg, bg, palette })
}
}
impl io::Read for DevTty {
@@ -104,6 +207,10 @@ impl Tty for NullTty {
ws_ypixel: 0,
}
}
fn get_theme(&self) -> Option<Theme> {
None
}
}
impl io::Read for NullTty {
@@ -127,3 +234,33 @@ impl AsFd for NullTty {
self.tx.as_fd()
}
}
#[cfg(test)]
mod tests {
use rgb::RGB8;
#[test]
fn parse_color() {
use super::parse_color as parse;
let color = Some(RGB8::new(0xaa, 0xbb, 0xcc));
assert_eq!(parse("aa11/bb22/cc33"), color);
assert_eq!(parse("aa11/bb22/cc33\x07"), color);
assert_eq!(parse("aa11/bb22/cc33\x1b\\"), color);
assert_eq!(parse("aa11/bb22/cc33.."), color);
assert_eq!(parse("aa1/bb2/cc3"), color);
assert_eq!(parse("aa1/bb2/cc3\x07"), color);
assert_eq!(parse("aa1/bb2/cc3\x1b\\"), color);
assert_eq!(parse("aa1/bb2/cc3.."), color);
assert_eq!(parse("aa/bb/cc"), color);
assert_eq!(parse("aa/bb/cc\x07"), color);
assert_eq!(parse("aa/bb/cc\x1b\\"), color);
assert_eq!(parse("aa/bb/cc.."), color);
assert_eq!(parse("aa11/bb22"), None);
assert_eq!(parse("xxxx/yyyy/zzzz"), None);
assert_eq!(parse("xxx/yyy/zzz"), None);
assert_eq!(parse("xx/yy/zz"), None);
assert_eq!(parse("foo"), None);
assert_eq!(parse(""), None);
}
}

6
tests/casts/full.cast Normal file
View File

@@ -0,0 +1,6 @@
{"version":2,"width":100,"height":50,"command":"/bin/bash","env":{"TERM":"xterm-256color","SHELL":"/bin/bash"},"theme":{"fg":"#000000","bg":"#ffffff","palette":"#241f31:#c01c28:#2ec27e:#f5c211:#1e78e4:#9841bb:#0ab9dc:#c0bfbc:#5e5c64:#ed333b:#57e389:#f8e45c:#51a1ff:#c061cb:#4fd2fd:#f6f5f4"}}
[0.000001, "o", "ż"]
[1.0, "o", "ółć"]
[2.3, "i", "\n"]
[5.600001, "r", "80x40"]
[10.5, "o", "\r\n"]

2
tests/casts/minimal.cast Normal file
View File

@@ -0,0 +1,2 @@
{"version":2,"width":100,"height":50}
[1.23, "o", "hello"]