From 7e676620eb5c09acb87340da7781d9c6fa1445ed Mon Sep 17 00:00:00 2001 From: Denis Drakhnia Date: Thu, 19 Oct 2023 15:38:13 +0300 Subject: [PATCH] protocol: color helpers query: colored output --- Cargo.lock | 1 + protocol/src/color.rs | 107 +++++++++++++++++++++++++++++++++++++++++ protocol/src/cursor.rs | 8 +-- protocol/src/lib.rs | 1 + query/Cargo.toml | 5 ++ query/src/cli.rs | 4 ++ query/src/main.rs | 45 ++++++++++++++++- 7 files changed, 163 insertions(+), 8 deletions(-) create mode 100644 protocol/src/color.rs diff --git a/Cargo.lock b/Cargo.lock index 403e8ee..a6d3bd0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -512,6 +512,7 @@ dependencies = [ "getopts", "serde", "serde_json", + "termion", "thiserror", "xash3d-protocol", ] diff --git a/protocol/src/color.rs b/protocol/src/color.rs new file mode 100644 index 0000000..0206514 --- /dev/null +++ b/protocol/src/color.rs @@ -0,0 +1,107 @@ +// SPDX-License-Identifier: GPL-3.0-only +// SPDX-FileCopyrightText: 2023 Denis Drakhnia + +use std::borrow::Cow; + +#[derive(Copy, Clone, Debug, PartialEq, Eq)] +pub enum Color { + Black, + Red, + Green, + Yellow, + Blue, + Cyan, + Magenta, + White, +} + +impl TryFrom<&str> for Color { + type Error = (); + + fn try_from(value: &str) -> Result { + Ok(match value { + "^0" => Self::Black, + "^1" => Self::Red, + "^2" => Self::Green, + "^3" => Self::Yellow, + "^4" => Self::Blue, + "^5" => Self::Cyan, + "^6" => Self::Magenta, + "^7" => Self::White, + _ => return Err(()), + }) + } +} + +#[inline] +pub fn is_color_code(s: &str) -> bool { + matches!(s.as_bytes(), [b'^', c, ..] if c.is_ascii_digit()) +} + +#[inline] +pub fn trim_start_color(s: &str) -> (&str, &str) { + let n = if is_color_code(s) { 2 } else { 0 }; + s.split_at(n) +} + +pub struct ColorIter<'a> { + inner: &'a str, +} + +impl<'a> ColorIter<'a> { + pub fn new(inner: &'a str) -> Self { + Self { inner } + } +} + +impl<'a> Iterator for ColorIter<'a> { + type Item = (&'a str, &'a str); + + fn next(&mut self) -> Option { + if !self.inner.is_empty() { + let i = self.inner[1..].find('^').map(|i| i + 1).unwrap_or(self.inner.len()); + let (head, tail) = self.inner.split_at(i); + let (color, text) = trim_start_color(head); + self.inner = tail; + Some((color, text)) + } else { + None + } + } +} + +pub fn trim_color(s: &str) -> Cow<'_, str> { + let (_, s) = trim_start_color(s); + if !s.chars().any(|c| c == '^') { + return Cow::Borrowed(s); + } + + let mut out = String::with_capacity(s.len()); + for (_, s) in ColorIter::new(s) { + out.push_str(s); + } + + Cow::Owned(out) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn trim_start_colors() { + assert_eq!(trim_start_color("foo^2bar"), ("", "foo^2bar")); + assert_eq!(trim_start_color("^foo^2bar"), ("", "^foo^2bar")); + assert_eq!(trim_start_color("^1foo^2bar"), ("^1", "foo^2bar")); + } + + #[test] + fn trim_colors() { + assert_eq!(trim_color("foo^2bar"), "foobar"); + assert_eq!(trim_color("^1foo^2bar^3"), "foobar"); + assert_eq!(trim_color("^1foo^2bar^3"), "foobar"); + assert_eq!(trim_color("^1foo^bar^3"), "foo^bar"); + assert_eq!(trim_color("^1foo^2bar^"), "foobar^"); + assert_eq!(trim_color("^foo^bar^"), "^foo^bar^"); + } +} diff --git a/protocol/src/cursor.rs b/protocol/src/cursor.rs index fc3cbeb..4087a54 100644 --- a/protocol/src/cursor.rs +++ b/protocol/src/cursor.rs @@ -8,7 +8,7 @@ use std::slice; use std::str; use super::types::Str; -use super::Error; +use super::{Error, color}; pub trait GetKeyValue<'a>: Sized { fn get_key_value(cur: &mut Cursor<'a>) -> Result; @@ -67,11 +67,7 @@ macro_rules! impl_get_value { fn get_key_value(cur: &mut Cursor<'a>) -> Result { let s = cur.get_key_value::<&str>()?; // HACK: special case for one asshole - let s = if s.len() > 2 && s.as_bytes()[0] == b'^' && s.as_bytes()[1].is_ascii_digit() { - &s[2..] - } else { - s - }; + let (_, s) = color::trim_start_color(s); s.parse().map_err(|_| Error::InvalidPacket) } })+ diff --git a/protocol/src/lib.rs b/protocol/src/lib.rs index dd082f0..1ee1bff 100644 --- a/protocol/src/lib.rs +++ b/protocol/src/lib.rs @@ -10,6 +10,7 @@ pub mod game; pub mod master; pub mod server; pub mod types; +pub mod color; pub use server_info::ServerInfo; diff --git a/query/Cargo.toml b/query/Cargo.toml index 88102e6..8d02c50 100644 --- a/query/Cargo.toml +++ b/query/Cargo.toml @@ -6,9 +6,14 @@ authors = ["Denis Drakhnia "] edition = "2021" rust-version = "1.56" +[features] +default = ["color"] +color = ["termion"] + [dependencies] thiserror = "1.0.49" getopts = "0.2.21" serde = { version = "1.0.188", features = ["derive"] } serde_json = "1.0.107" +termion = { version = "2", optional = true } xash3d-protocol = { path = "../protocol", version = "0.1.0" } diff --git a/query/src/cli.rs b/query/src/cli.rs index 361831b..b97e0aa 100644 --- a/query/src/cli.rs +++ b/query/src/cli.rs @@ -21,6 +21,7 @@ pub struct Cli { pub protocol: Vec, pub json: bool, pub debug: bool, + pub force_color: bool, } impl Default for Cli { @@ -36,6 +37,7 @@ impl Default for Cli { protocol: vec![xash3d_protocol::VERSION, xash3d_protocol::VERSION - 1], json: false, debug: false, + force_color: false, } } } @@ -91,6 +93,7 @@ pub fn parse() -> Cli { opts.optopt("p", "protocol", &help, "VERSION"); opts.optflag("j", "json", "output JSON"); opts.optflag("d", "debug", "output debug"); + opts.optflag("F", "force-color", "force colored output"); let matches = match opts.parse(&args[1..]) { Ok(m) => m, @@ -160,6 +163,7 @@ pub fn parse() -> Cli { cli.json = matches.opt_present("json"); cli.debug = matches.opt_present("debug"); + cli.force_color = matches.opt_present("force-color"); cli.args = matches.free; cli diff --git a/query/src/main.rs b/query/src/main.rs index 4aec381..210ed12 100644 --- a/query/src/main.rs +++ b/query/src/main.rs @@ -16,7 +16,7 @@ use std::time::{Duration, Instant}; use serde::Serialize; use thiserror::Error; use xash3d_protocol::types::Str; -use xash3d_protocol::{filter, game, master, server, Error as ProtocolError}; +use xash3d_protocol::{color, filter, game, master, server, Error as ProtocolError}; use crate::cli::Cli; @@ -145,6 +145,47 @@ struct ListResult<'a> { servers: &'a [&'a str], } +struct Colored<'a> { + inner: &'a str, + forced: bool, +} + +impl<'a> Colored<'a> { + fn new(s: &'a str, forced: bool) -> Self { + Self { inner: s, forced } + } +} + +impl fmt::Display for Colored<'_> { + fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result { + #[cfg(feature = "color")] + if self.forced || termion::is_tty(&io::stdout()) { + use termion::color::*; + + for (color, text) in color::ColorIter::new(self.inner) { + match color::Color::try_from(color) { + Ok(color::Color::Black) => write!(fmt, "{}", Fg(Black))?, + Ok(color::Color::Red) => write!(fmt, "{}", Fg(Red))?, + Ok(color::Color::Green) => write!(fmt, "{}", Fg(Green))?, + Ok(color::Color::Yellow) => write!(fmt, "{}", Fg(Yellow))?, + Ok(color::Color::Blue) => write!(fmt, "{}", Fg(Blue))?, + Ok(color::Color::Cyan) => write!(fmt, "{}", Fg(Cyan))?, + Ok(color::Color::Magenta) => write!(fmt, "{}", Fg(Magenta))?, + Ok(color::Color::White) => write!(fmt, "{}", Fg(White))?, + _ => {} + } + write!(fmt, "{}", text)?; + } + return write!(fmt, "{}", Fg(Reset)); + } + + for (_, text) in color::ColorIter::new(self.inner) { + write!(fmt, "{}", text)?; + } + Ok(()) + } +} + enum Message { Servers(Vec), ServerResult(ServerResult), @@ -324,7 +365,7 @@ fn query_server_info(cli: &Cli, servers: &[String]) -> Result<(), Error> { ServerResultKind::Ok { info } => { p! { type: "ok", - host: info.host, + host: Colored::new(&info.host, cli.force_color), gamedir: info.gamedir, map: info.map, protocol: info.protocol,