Martin von Zweigbergkformatter: reset colors on early drop46a8afe3 years agoHistory
1038 lines · 35.6 KB · Rust
RawDownload
// Copyright 2020 The Jujutsu Authors//// Licensed under the Apache License, Version 2.0 (the "License");// you may not use this file except in compliance with the License.// You may obtain a copy of the License at//// https://www.apache.org/licenses/LICENSE-2.0//// Unless required by applicable law or agreed to in writing, software// distributed under the License is distributed on an "AS IS" BASIS,// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.// See the License for the specific language governing permissions and// limitations under the License. use std::borrow::BorrowMut;use std::collections::HashMap;use std::io::{Error, Write};use std::ops::Range;use std::sync::Arc;use std::{fmt, io, mem}; use crossterm::queue;use crossterm::style::{Attribute, Color, SetAttribute, SetBackgroundColor, SetForegroundColor};use itertools::Itertools; // Lets the caller label strings and translates the labels to colorspub trait Formatter: Write {    fn write_str(&mut self, text: &str) -> io::Result<()> {        self.write_all(text.as_bytes())    }     /// Returns the backing `Write`. This is useful for writing data that is    /// already formatted, such as in the graphical log.    fn raw(&mut self) -> &mut dyn Write;     fn push_label(&mut self, label: &str) -> io::Result<()>;     fn pop_label(&mut self) -> io::Result<()>;} impl dyn Formatter + '_ {    pub fn labeled<S: AsRef<str>>(&mut self, label: S) -> LabeledWriter<&mut Self, S> {        LabeledWriter {            formatter: self,            label,        }    }     pub fn with_label(        &mut self,        label: &str,        write_inner: impl FnOnce(&mut dyn Formatter) -> io::Result<()>,    ) -> io::Result<()> {        self.push_label(label)?;        // Call `pop_label()` whether or not `write_inner()` fails, but don't let        // its error replace the one from `write_inner()`.        write_inner(self).and(self.pop_label())    }} /// `Formatter` wrapper to write a labeled message with `write!()` or/// `writeln!()`.pub struct LabeledWriter<T, S> {    formatter: T,    label: S,} impl<T, S> LabeledWriter<T, S> {    pub fn new(formatter: T, label: S) -> Self {        LabeledWriter { formatter, label }    }} impl<'a, T, S> LabeledWriter<T, S>where    T: BorrowMut<dyn Formatter + 'a>,    S: AsRef<str>,{    pub fn write_fmt(&mut self, args: fmt::Arguments<'_>) -> io::Result<()> {        self.formatter            .borrow_mut()            .with_label(self.label.as_ref(), |formatter| formatter.write_fmt(args))    }} type Rules = Vec<(Vec<String>, Style)>; /// Creates `Formatter` instances with preconfigured parameters.#[derive(Clone, Debug)]pub struct FormatterFactory {    kind: FormatterFactoryKind,} #[derive(Clone, Debug)]enum FormatterFactoryKind {    PlainText,    Sanitized,    Color { rules: Arc<Rules> },} impl FormatterFactory {    pub fn prepare(        config: &config::Config,        color: bool,        sanitized: bool,    ) -> Result<Self, config::ConfigError> {        let kind = if color {            let rules = Arc::new(rules_from_config(config)?);            FormatterFactoryKind::Color { rules }        } else if sanitized {            FormatterFactoryKind::Sanitized        } else {            FormatterFactoryKind::PlainText        };        Ok(FormatterFactory { kind })    }     pub fn new_formatter<'output, W: Write + 'output>(        &self,        output: W,    ) -> Box<dyn Formatter + 'output> {        match &self.kind {            FormatterFactoryKind::PlainText => Box::new(PlainTextFormatter::new(output)),            FormatterFactoryKind::Sanitized => Box::new(SanitizingFormatter::new(output)),            FormatterFactoryKind::Color { rules } => {                Box::new(ColorFormatter::new(output, rules.clone()))            }        }    }} pub struct PlainTextFormatter<W> {    output: W,} impl<W> PlainTextFormatter<W> {    pub fn new(output: W) -> PlainTextFormatter<W> {        Self { output }    }} impl<W: Write> Write for PlainTextFormatter<W> {    fn write(&mut self, data: &[u8]) -> Result<usize, Error> {        self.output.write(data)    }     fn flush(&mut self) -> Result<(), Error> {        self.output.flush()    }} impl<W: Write> Formatter for PlainTextFormatter<W> {    fn raw(&mut self) -> &mut dyn Write {        &mut self.output    }     fn push_label(&mut self, _label: &str) -> io::Result<()> {        Ok(())    }     fn pop_label(&mut self) -> io::Result<()> {        Ok(())    }} pub struct SanitizingFormatter<W> {    output: W,} impl<W> SanitizingFormatter<W> {    pub fn new(output: W) -> SanitizingFormatter<W> {        Self { output }    }} impl<W: Write> Write for SanitizingFormatter<W> {    fn write(&mut self, data: &[u8]) -> Result<usize, Error> {        write_sanitized(&mut self.output, data)?;        Ok(data.len())    }     fn flush(&mut self) -> Result<(), Error> {        self.output.flush()    }} impl<W: Write> Formatter for SanitizingFormatter<W> {    fn raw(&mut self) -> &mut dyn Write {        &mut self.output    }     fn push_label(&mut self, _label: &str) -> io::Result<()> {        Ok(())    }     fn pop_label(&mut self) -> io::Result<()> {        Ok(())    }} #[derive(Clone, Debug, Default, PartialEq, Eq)]pub struct Style {    pub fg_color: Option<Color>,    pub bg_color: Option<Color>,    pub bold: Option<bool>,    pub underlined: Option<bool>,} impl Style {    fn merge(&mut self, other: &Style) {        self.fg_color = other.fg_color.or(self.fg_color);        self.bg_color = other.bg_color.or(self.bg_color);        self.bold = other.bold.or(self.bold);        self.underlined = other.underlined.or(self.underlined);    }} #[derive(Clone, Debug)]pub struct ColorFormatter<W: Write> {    output: W,    rules: Arc<Rules>,    /// The stack of currently applied labels. These determine the desired    /// style.    labels: Vec<String>,    cached_styles: HashMap<Vec<String>, Style>,    /// The style we last wrote to the output.    current_style: Style,} impl<W: Write> ColorFormatter<W> {    pub fn new(output: W, rules: Arc<Rules>) -> ColorFormatter<W> {        ColorFormatter {            output,            rules,            labels: vec![],            cached_styles: HashMap::new(),            current_style: Style::default(),        }    }     pub fn for_config(output: W, config: &config::Config) -> Result<Self, config::ConfigError> {        let rules = rules_from_config(config)?;        Ok(Self::new(output, Arc::new(rules)))    }     fn requested_style(&mut self) -> Style {        if let Some(cached) = self.cached_styles.get(&self.labels) {            cached.clone()        } else {            // We use the reverse list of matched indices as a measure of how well the rule            // matches the actual labels. For example, for rule "a d" and the actual labels            // "a b c d", we'll get [3,0]. We compare them by Rust's default Vec comparison.            // That means "a d" will trump both rule "d" (priority [3]) and rule            // "a b c" (priority [2,1,0]).            let mut matched_styles = vec![];            for (labels, style) in self.rules.as_ref() {                let mut labels_iter = self.labels.iter().enumerate();                // The indexes in the current label stack that match the required label.                let mut matched_indices = vec![];                for required_label in labels {                    for (label_index, label) in &mut labels_iter {                        if label == required_label {                            matched_indices.push(label_index);                            break;                        }                    }                }                if matched_indices.len() == labels.len() {                    matched_indices.reverse();                    matched_styles.push((style, matched_indices));                }            }            matched_styles.sort_by_key(|(_, indices)| indices.clone());             let mut style = Style::default();            for (matched_style, _) in matched_styles {                style.merge(matched_style);            }            self.cached_styles                .insert(self.labels.clone(), style.clone());            style        }    }     fn write_new_style(&mut self) -> io::Result<()> {        let new_style = self.requested_style();        if new_style != self.current_style {            if new_style.bold != self.current_style.bold {                if new_style.bold.unwrap_or_default() {                    queue!(self.output, SetAttribute(Attribute::Bold))?;                } else {                    // NoBold results in double underlining on some terminals, so we use reset                    // instead. However, that resets other attributes as well, so we reset                    // our record of the current style so we re-apply the other attributes                    // below.                    queue!(self.output, SetAttribute(Attribute::Reset))?;                    self.current_style = Style::default();                }            }            if new_style.underlined != self.current_style.underlined {                if new_style.underlined.unwrap_or_default() {                    queue!(self.output, SetAttribute(Attribute::Underlined))?;                } else {                    queue!(self.output, SetAttribute(Attribute::NoUnderline))?;                }            }            if new_style.fg_color != self.current_style.fg_color {                queue!(                    self.output,                    SetForegroundColor(new_style.fg_color.unwrap_or(Color::Reset))                )?;            }            if new_style.bg_color != self.current_style.bg_color {                queue!(                    self.output,                    SetBackgroundColor(new_style.bg_color.unwrap_or(Color::Reset))                )?;            }            self.current_style = new_style;        }        Ok(())    }} fn rules_from_config(config: &config::Config) -> Result<Rules, config::ConfigError> {    let mut result = vec![];    let table = config.get_table("colors")?;    for (key, value) in table {        let labels = key            .split_whitespace()            .map(ToString::to_string)            .collect_vec();        match value.kind {            config::ValueKind::String(color_name) => {                let style = Style {                    fg_color: Some(color_for_name(&color_name)?),                    bg_color: None,                    bold: None,                    underlined: None,                };                result.push((labels, style));            }            config::ValueKind::Table(style_table) => {                let mut style = Style::default();                if let Some(value) = style_table.get("fg") {                    if let config::ValueKind::String(color_name) = &value.kind {                        style.fg_color = Some(color_for_name(color_name)?);                    }                }                if let Some(value) = style_table.get("bg") {                    if let config::ValueKind::String(color_name) = &value.kind {                        style.bg_color = Some(color_for_name(color_name)?);                    }                }                if let Some(value) = style_table.get("bold") {                    if let config::ValueKind::Boolean(value) = &value.kind {                        style.bold = Some(*value);                    }                }                if let Some(value) = style_table.get("underline") {                    if let config::ValueKind::Boolean(value) = &value.kind {                        style.underlined = Some(*value);                    }                }                result.push((labels, style));            }            _ => {}        }    }    Ok(result)} fn color_for_name(color_name: &str) -> Result<Color, config::ConfigError> {    match color_name {        "default" => Ok(Color::Reset),        "black" => Ok(Color::Black),        "red" => Ok(Color::DarkRed),        "green" => Ok(Color::DarkGreen),        "yellow" => Ok(Color::DarkYellow),        "blue" => Ok(Color::DarkBlue),        "magenta" => Ok(Color::DarkMagenta),        "cyan" => Ok(Color::DarkCyan),        "white" => Ok(Color::Grey),        "bright black" => Ok(Color::DarkGrey),        "bright red" => Ok(Color::Red),        "bright green" => Ok(Color::Green),        "bright yellow" => Ok(Color::Yellow),        "bright blue" => Ok(Color::Blue),        "bright magenta" => Ok(Color::Magenta),        "bright cyan" => Ok(Color::Cyan),        "bright white" => Ok(Color::White),        _ => Err(config::ConfigError::Message(format!(            "invalid color: {color_name}"        ))),    }} impl<W: Write> Write for ColorFormatter<W> {    fn write(&mut self, data: &[u8]) -> Result<usize, Error> {        /*        We clear the current style at the end of each line, and then we re-apply the style        after the newline. There are several reasons for this:          * We can more easily skip styling a trailing blank line, which other           internal code then can correctly detect as having a trailing           newline.          * Some tools (like `less -R`) add an extra newline if the final           character is not a newline (e.g. if there's a color reset after           it), which led to an annoying blank line after the diff summary in           e.g. `jj status`.          * Since each line is styled independently, you get all the necessary           escapes even when grepping through the output.          * Some terminals extend background color to the end of the terminal           (i.e. past the newline character), which is probably not what the           user wanted.          * Some tools (like `less -R`) get confused and lose coloring of lines           after a newline.         */        for line in data.split_inclusive(|b| *b == b'\n') {            if line.ends_with(b"\n") {                self.write_new_style()?;                write_sanitized(&mut self.output, &line[..line.len() - 1])?;                let labels = mem::take(&mut self.labels);                self.write_new_style()?;                self.output.write_all(b"\n")?;                self.labels = labels;            } else {                self.write_new_style()?;                write_sanitized(&mut self.output, line)?;            }        }        Ok(data.len())    }     fn flush(&mut self) -> Result<(), Error> {        self.output.flush()    }} impl<W: Write> Formatter for ColorFormatter<W> {    fn raw(&mut self) -> &mut dyn Write {        &mut self.output    }     fn push_label(&mut self, label: &str) -> io::Result<()> {        self.labels.push(label.to_owned());        Ok(())    }     fn pop_label(&mut self) -> io::Result<()> {        self.labels.pop();        if self.labels.is_empty() {            self.write_new_style()?        }        Ok(())    }} impl<W: Write> Drop for ColorFormatter<W> {    fn drop(&mut self) {        // If a `ColorFormatter` was dropped without popping all labels first (perhaps        // because of an error), let's still try to reset any currently active style.        self.labels.clear();        self.write_new_style().ok();    }} /// Like buffered formatter, but records `push`/`pop_label()` calls.////// This allows you to manipulate the recorded data without losing labels./// The recorded data and labels can be written to another formatter. If/// the destination formatter has already been labeled, the recorded labels/// will be stacked on top of the existing labels, and the subsequent data/// may be colorized differently.#[derive(Debug, Default)]pub struct FormatRecorder {    data: Vec<u8>,    label_ops: Vec<(usize, LabelOp)>,} #[derive(Clone, Debug, Eq, PartialEq)]enum LabelOp {    PushLabel(String),    PopLabel,} impl FormatRecorder {    pub fn new() -> Self {        FormatRecorder::default()    }     pub fn data(&self) -> &[u8] {        &self.data    }     fn push_label_op(&mut self, op: LabelOp) {

Lines 1–500 of 1038.