protocol: color helpers

query: colored output

Cargo.lock

@@ -512,6 +512,7 @@ dependencies = [
  "getopts",
  "serde",
  "serde_json",
+ "termion",
  "thiserror",
  "xash3d-protocol",
 ]

protocol/src/color.rs

@@ -0,0 +1,107 @@
+// SPDX-License-Identifier: GPL-3.0-only
+// SPDX-FileCopyrightText: 2023 Denis Drakhnia <numas13@gmail.com>
+
+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<Self, Self::Error> {
+        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<Self::Item> {
+        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^");
+    }
+}

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<Self, Error>;
@@ -67,11 +67,7 @@ macro_rules! impl_get_value {
             fn get_key_value(cur: &mut Cursor<'a>) -> Result<Self, Error> {
                 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)
             }
         })+

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;
 

query/Cargo.toml

@@ -6,9 +6,14 @@ authors = ["Denis Drakhnia <numas13@gmail.com>"]
 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" }

query/src/cli.rs

@@ -21,6 +21,7 @@ pub struct Cli {
     pub protocol: Vec<u8>,
     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

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<String>),
     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,