rhg: Add support for colored output

The same "label" system is used as in Python code

Differential Revision: https://phab.mercurial-scm.org/D12167

rust/rhg/src/color.rs

@@ -0,0 +1,255 b''
+use crate::ui::formatted;
+use crate::ui::plain;
+use format_bytes::write_bytes;
+use hg::config::Config;
+use hg::config::ConfigOrigin;
+use hg::errors::HgError;
+use std::collections::HashMap;
+
+pub type Effect = u32;
+
+pub type EffectsMap = HashMap<Vec<u8>, Vec<Effect>>;
+
+macro_rules! effects {
+    ($( $name: ident: $value: expr ,)+) => {
+
+        #[allow(non_upper_case_globals)]
+        mod effects {
+            $(
+                pub const $name: super::Effect = $value;
+            )+
+        }
+
+        fn effect(name: &[u8]) -> Option<Effect> {
+            $(
+                if name == stringify!($name).as_bytes() {
+                    Some(effects::$name)
+                } else
+            )+
+            {
+                None
+            }
+        }
+    };
+}
+
+effects! {
+    none: 0,
+    black: 30,
+    red: 31,
+    green: 32,
+    yellow: 33,
+    blue: 34,
+    magenta: 35,
+    cyan: 36,
+    white: 37,
+    bold: 1,
+    italic: 3,
+    underline: 4,
+    inverse: 7,
+    dim: 2,
+    black_background: 40,
+    red_background: 41,
+    green_background: 42,
+    yellow_background: 43,
+    blue_background: 44,
+    purple_background: 45,
+    cyan_background: 46,
+    white_background: 47,
+}
+
+macro_rules! default_styles {
+    ($( $key: expr => [$($value: expr),*],)+) => {
+        fn default_styles() -> EffectsMap {
+            use effects::*;
+            let mut map = HashMap::new();
+            $(
+                map.insert($key[..].to_owned(), vec![$( $value ),*]);
+            )+
+            map
+        }
+    };
+}
+
+default_styles! {
+    b"grep.match" => [red, bold],
+    b"grep.linenumber" => [green],
+    b"grep.rev" => [blue],
+    b"grep.sep" => [cyan],
+    b"grep.filename" => [magenta],
+    b"grep.user" => [magenta],
+    b"grep.date" => [magenta],
+    b"grep.inserted" => [green, bold],
+    b"grep.deleted" => [red, bold],
+    b"bookmarks.active" => [green],
+    b"branches.active" => [none],
+    b"branches.closed" => [black, bold],
+    b"branches.current" => [green],
+    b"branches.inactive" => [none],
+    b"diff.changed" => [white],
+    b"diff.deleted" => [red],
+    b"diff.deleted.changed" => [red, bold, underline],
+    b"diff.deleted.unchanged" => [red],
+    b"diff.diffline" => [bold],
+    b"diff.extended" => [cyan, bold],
+    b"diff.file_a" => [red, bold],
+    b"diff.file_b" => [green, bold],
+    b"diff.hunk" => [magenta],
+    b"diff.inserted" => [green],
+    b"diff.inserted.changed" => [green, bold, underline],
+    b"diff.inserted.unchanged" => [green],
+    b"diff.tab" => [],
+    b"diff.trailingwhitespace" => [bold, red_background],
+    b"changeset.public" => [],
+    b"changeset.draft" => [],
+    b"changeset.secret" => [],
+    b"diffstat.deleted" => [red],
+    b"diffstat.inserted" => [green],
+    b"formatvariant.name.mismatchconfig" => [red],
+    b"formatvariant.name.mismatchdefault" => [yellow],
+    b"formatvariant.name.uptodate" => [green],
+    b"formatvariant.repo.mismatchconfig" => [red],
+    b"formatvariant.repo.mismatchdefault" => [yellow],
+    b"formatvariant.repo.uptodate" => [green],
+    b"formatvariant.config.special" => [yellow],
+    b"formatvariant.config.default" => [green],
+    b"formatvariant.default" => [],
+    b"histedit.remaining" => [red, bold],
+    b"ui.addremove.added" => [green],
+    b"ui.addremove.removed" => [red],
+    b"ui.error" => [red],
+    b"ui.prompt" => [yellow],
+    b"log.changeset" => [yellow],
+    b"patchbomb.finalsummary" => [],
+    b"patchbomb.from" => [magenta],
+    b"patchbomb.to" => [cyan],
+    b"patchbomb.subject" => [green],
+    b"patchbomb.diffstats" => [],
+    b"rebase.rebased" => [blue],
+    b"rebase.remaining" => [red, bold],
+    b"resolve.resolved" => [green, bold],
+    b"resolve.unresolved" => [red, bold],
+    b"shelve.age" => [cyan],
+    b"shelve.newest" => [green, bold],
+    b"shelve.name" => [blue, bold],
+    b"status.added" => [green, bold],
+    b"status.clean" => [none],
+    b"status.copied" => [none],
+    b"status.deleted" => [cyan, bold, underline],
+    b"status.ignored" => [black, bold],
+    b"status.modified" => [blue, bold],
+    b"status.removed" => [red, bold],
+    b"status.unknown" => [magenta, bold, underline],
+    b"tags.normal" => [green],
+    b"tags.local" => [black, bold],
+    b"upgrade-repo.requirement.preserved" => [cyan],
+    b"upgrade-repo.requirement.added" => [green],
+    b"upgrade-repo.requirement.removed" => [red],
+}
+
+fn parse_effect(config_key: &[u8], effect_name: &[u8]) -> Option<Effect> {
+    let found = effect(effect_name);
+    if found.is_none() {
+        // TODO: have some API for warnings
+        // TODO: handle IO errors during warnings
+        let stderr = std::io::stderr();
+        let _ = write_bytes!(
+            &mut stderr.lock(),
+            b"ignoring unknown color/effect '{}' \
+              (configured in color.{})\n",
+            effect_name,
+            config_key,
+        );
+    }
+    found
+}
+
+fn effects_from_config(config: &Config) -> EffectsMap {
+    let mut styles = default_styles();
+    for (key, _value) in config.iter_section(b"color") {
+        if !key.contains(&b'.')
+            || key.starts_with(b"color.")
+            || key.starts_with(b"terminfo.")
+        {
+            continue;
+        }
+        // `unwrap` shouldn’t panic since we just got this key from
+        // iteration
+        let list = config.get_list(b"color", key).unwrap();
+        let parsed = list
+            .iter()
+            .filter_map(|name| parse_effect(key, name))
+            .collect();
+        styles.insert(key.to_owned(), parsed);
+    }
+    styles
+}
+
+enum ColorMode {
+    // TODO: support other modes
+    Ansi,
+}
+
+impl ColorMode {
+    // Similar to _modesetup in mercurial/color.py
+    fn get(config: &Config) -> Result<Option<Self>, HgError> {
+        if plain(Some("color")) {
+            return Ok(None);
+        }
+        let enabled_default = b"auto";
+        // `origin` is only used when `!auto`, so its default doesn’t matter
+        let (enabled, origin) = config
+            .get_with_origin(b"ui", b"color")
+            .unwrap_or((enabled_default, &ConfigOrigin::CommandLineColor));
+        if enabled == b"debug" {
+            return Err(HgError::unsupported("debug color mode"));
+        }
+        let auto = enabled == b"auto";
+        let always;
+        if !auto {
+            let enabled_bool = config.get_bool(b"ui", b"color")?;
+            if !enabled_bool {
+                return Ok(None);
+            }
+            always = enabled == b"always"
+                || *origin == ConfigOrigin::CommandLineColor
+        } else {
+            always = false
+        };
+        let formatted = always
+            || (std::env::var_os("TERM").unwrap_or_default() != "dumb"
+                && formatted(config)?);
+
+        let mode_default = b"auto";
+        let mode = config.get(b"color", b"mode").unwrap_or(mode_default);
+
+        if formatted {
+            match mode {
+                b"ansi" | b"auto" => Ok(Some(ColorMode::Ansi)),
+                // TODO: support other modes
+                _ => Err(HgError::UnsupportedFeature(format!(
+                    "color mode {}",
+                    String::from_utf8_lossy(mode)
+                ))),
+            }
+        } else {
+            Ok(None)
+        }
+    }
+}
+
+pub struct ColorConfig {
+    pub styles: EffectsMap,
+}
+
+impl ColorConfig {
+    // Similar to _modesetup in mercurial/color.py
+    pub fn new(config: &Config) -> Result<Option<Self>, HgError> {
+        Ok(match ColorMode::get(config)? {
+            None => None,
+            Some(ColorMode::Ansi) => Some(ColorConfig {
+                styles: effects_from_config(config),
+            }),
+        })
+    }
+}

rust/Cargo.lock

@@ -876,6 +876,7 b' dependencies = ['
 name = "rhg"
 version = "0.1.0"
 dependencies = [
+ "atty",
  "chrono",
  "clap",
  "derive_more",

rust/hg-core/src/config.rs

@@ -13,4 +13,4 b' mod config;'
 mod layer;
 mod values;
 pub use config::{Config, ConfigSource, ConfigValueParseError};
-pub use layer::{ConfigError, ConfigParseError};
+pub use layer::{ConfigError, ConfigOrigin, ConfigParseError};

rust/hg-core/src/config/config.rs

@@ -398,6 +398,16 b' impl Config {'
             .map(|(_, value)| value.bytes.as_ref())
     }
 
+    /// Returns the raw value bytes of the first one found, or `None`.
+    pub fn get_with_origin(
+        &self,
+        section: &[u8],
+        item: &[u8],
+    ) -> Option<(&[u8], &ConfigOrigin)> {
+        self.get_inner(section, item)
+            .map(|(layer, value)| (value.bytes.as_ref(), &layer.origin))
+    }
+
     /// Returns the layer and the value of the first one found, or `None`.
     fn get_inner(
         &self,

rust/hg-core/src/config/layer.rs

@@ -295,7 +295,7 b' pub struct ConfigValue {'
     pub line: Option<usize>,
 }
 
-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, PartialEq, Eq)]
 pub enum ConfigOrigin {
     /// From a configuration file
     File(PathBuf),

rust/rhg/Cargo.toml

@@ -8,6 +8,7 b' authors = ['
 edition = "2018"
 
 [dependencies]
+atty = "0.2"
 hg-core = { path = "../hg-core"}
 chrono = "0.4.19"
 clap = "2.33.1"

rust/rhg/src/main.rs

@@ -17,6 +17,7 b' use std::path::PathBuf;'
 use std::process::Command;
 
 mod blackbox;
+mod color;
 mod error;
 mod ui;
 pub mod utils {

rust/rhg/src/ui.rs

@@ -1,4 +1,7 b''
+use crate::color::ColorConfig;
+use crate::color::Effect;
 use format_bytes::format_bytes;
+use format_bytes::write_bytes;
 use hg::config::Config;
 use hg::errors::HgError;
 use hg::utils::files::get_bytes_from_os_string;
@@ -7,10 +10,10 b' use std::env;'
 use std::io;
 use std::io::{ErrorKind, Write};
 
-#[derive(Debug)]
 pub struct Ui {
     stdout: std::io::Stdout,
     stderr: std::io::Stderr,
+    colors: Option<ColorConfig>,
 }
 
 /// The kind of user interface error
@@ -23,20 +26,26 b' pub enum UiError {'
 
 /// The commandline user interface
 impl Ui {
-    pub fn new(_config: &Config) -> Result<Self, HgError> {
+    pub fn new(config: &Config) -> Result<Self, HgError> {
         Ok(Ui {
+            // If using something else, also adapt `isatty()` below.
             stdout: std::io::stdout(),
+
             stderr: std::io::stderr(),
+            colors: ColorConfig::new(config)?,
         })
     }
 
     /// Default to no color if color configuration errors.
     ///
     /// Useful when we’re already handling another error.
-    pub fn new_infallible(_config: &Config) -> Self {
+    pub fn new_infallible(config: &Config) -> Self {
         Ui {
+            // If using something else, also adapt `isatty()` below.
             stdout: std::io::stdout(),
+
             stderr: std::io::stderr(),
+            colors: ColorConfig::new(config).unwrap_or(None),
         }
     }
 
@@ -48,6 +57,11 b' impl Ui {'
 
     /// Write bytes to stdout
     pub fn write_stdout(&self, bytes: &[u8]) -> Result<(), UiError> {
+        // Hack to silence "unused" warnings
+        if false {
+            return self.write_stdout_labelled(bytes, "");
+        }
+
         let mut stdout = self.stdout.lock();
 
         stdout.write_all(bytes).or_else(handle_stdout_error)?;
@@ -64,6 +78,61 b' impl Ui {'
         stderr.flush().or_else(handle_stderr_error)
     }
 
+    /// Write bytes to stdout with the given label
+    ///
+    /// Like the optional `label` parameter in `mercurial/ui.py`,
+    /// this label influences the color used for this output.
+    pub fn write_stdout_labelled(
+        &self,
+        bytes: &[u8],
+        label: &str,
+    ) -> Result<(), UiError> {
+        if let Some(colors) = &self.colors {
+            if let Some(effects) = colors.styles.get(label.as_bytes()) {
+                if !effects.is_empty() {
+                    return self
+                        .write_stdout_with_effects(bytes, effects)
+                        .or_else(handle_stdout_error);
+                }
+            }
+        }
+        self.write_stdout(bytes)
+    }
+
+    fn write_stdout_with_effects(
+        &self,
+        bytes: &[u8],
+        effects: &[Effect],
+    ) -> io::Result<()> {
+        let stdout = &mut self.stdout.lock();
+        let mut write_line = |line: &[u8], first: bool| {
+            // `line` does not include the newline delimiter
+            if !first {
+                stdout.write_all(b"\n")?;
+            }
+            if line.is_empty() {
+                return Ok(());
+            }
+            /// 0x1B == 27 == 0o33
+            const ASCII_ESCAPE: &[u8] = b"\x1b";
+            write_bytes!(stdout, b"{}[0", ASCII_ESCAPE)?;
+            for effect in effects {
+                write_bytes!(stdout, b";{}", effect)?;
+            }
+            write_bytes!(stdout, b"m")?;
+            stdout.write_all(line)?;
+            write_bytes!(stdout, b"{}[0m", ASCII_ESCAPE)
+        };
+        let mut lines = bytes.split(|&byte| byte == b'\n');
+        if let Some(first) = lines.next() {
+            write_line(first, true)?;
+            for line in lines {
+                write_line(line, false)?
+            }
+        }
+        stdout.flush()
+    }
+
     /// Return whether plain mode is active.
     ///
     /// Plain mode means that all configuration variables which affect
@@ -83,7 +152,7 b' impl Ui {'
     }
 }
 
-fn plain(opt_feature: Option<&str>) -> bool {
+pub fn plain(opt_feature: Option<&str>) -> bool {
     if let Some(except) = env::var_os("HGPLAINEXCEPT") {
         opt_feature.map_or(true, |feature| {
             get_bytes_from_os_string(except)
@@ -154,3 +223,23 b' pub fn utf8_to_local(s: &str) -> Cow<[u8'
     let bytes = s.as_bytes();
     Cow::Borrowed(bytes)
 }
+
+/// Should formatted output be used?
+///
+/// Note: rhg does not have the formatter mechanism yet,
+/// but this is also used when deciding whether to use color.
+pub fn formatted(config: &Config) -> Result<bool, HgError> {
+    if let Some(formatted) = config.get_option(b"ui", b"formatted")? {
+        Ok(formatted)
+    } else {
+        isatty(config)
+    }
+}
+
+fn isatty(config: &Config) -> Result<bool, HgError> {
+    Ok(if config.get_bool(b"ui", b"nontty")? {
+        false
+    } else {
+        atty::is(atty::Stream::Stdout)
+    })
+}

tests/test-status-color.t

@@ -313,6 +313,7 b' test unknown color'
   ignoring unknown color/effect 'periwinkle' (configured in color.status.modified)
   ignoring unknown color/effect 'periwinkle' (configured in color.status.modified)
   ignoring unknown color/effect 'periwinkle' (configured in color.status.modified)
+  ignoring unknown color/effect 'periwinkle' (configured in color.status.modified) (rhg !)
   M modified
   \x1b[0;32;1mA \x1b[0m\x1b[0;32;1madded\x1b[0m (esc)
   \x1b[0;32;1mA \x1b[0m\x1b[0;32;1mcopied\x1b[0m (esc)