bat/src/main.rs

// `error_chain!` can recurse deeply
#![recursion_limit = "1024"]

#[macro_use]
extern crate error_chain;

#[macro_use]
extern crate clap;

#[macro_use]
extern crate lazy_static;

extern crate ansi_term;
extern crate atty;
extern crate console;
extern crate directories;
extern crate git2;
extern crate syntect;

mod assets;
mod diff;
mod printer;
mod terminal;

use std::collections::{HashMap, HashSet};
use std::env;
use std::fs::{self, File};
use std::io::{self, BufRead, BufReader, Write};
use std::process::{self, Child, Command, Stdio};
use std::str::FromStr;

#[cfg(unix)]
use std::os::unix::fs::FileTypeExt;

use ansi_term::Colour::{Fixed, Green, Red, White, Yellow};
use ansi_term::Style;
use atty::Stream;
use clap::{App, AppSettings, Arg, ArgGroup, SubCommand};

use syntect::easy::HighlightLines;
use syntect::highlighting::Theme;
use syntect::parsing::SyntaxSet;

use assets::{config_dir, syntax_set_path, theme_set_path, HighlightingAssets};
use diff::get_git_diff;
use printer::Printer;

mod errors {
    error_chain! {
        foreign_links {
            Clap(::clap::Error);
            Io(::std::io::Error);
        }

        errors {
            NoCorrectStylesSpecified {
                description("no correct styles specified")
            }

            UnknownStyleName(name: String) {
                description("unknown style name")
                display("unknown style name: '{}'", name)
            }
        }
    }
}

use errors::*;

#[derive(Debug, Eq, PartialEq, Copy, Clone, Hash)]
pub enum OutputComponent {
    Changes,
    Grid,
    Header,
    Numbers,
    Full,
    Plain,
}

impl OutputComponent {
    fn components(&self) -> &'static [OutputComponent] {
        match *self {
            OutputComponent::Changes => &[OutputComponent::Changes],
            OutputComponent::Grid => &[OutputComponent::Grid],
            OutputComponent::Header => &[OutputComponent::Header],
            OutputComponent::Numbers => &[OutputComponent::Numbers],
            OutputComponent::Full => &[
                OutputComponent::Changes,
                OutputComponent::Grid,
                OutputComponent::Header,
                OutputComponent::Numbers,
            ],
            OutputComponent::Plain => &[],
        }
    }
}

impl FromStr for OutputComponent {
    type Err = Error;

    fn from_str(s: &str) -> Result<Self> {
        match s {
            "changes" => Ok(OutputComponent::Changes),
            "grid" => Ok(OutputComponent::Grid),
            "header" => Ok(OutputComponent::Header),
            "numbers" => Ok(OutputComponent::Numbers),
            "full" => Ok(OutputComponent::Full),
            "plain" => Ok(OutputComponent::Plain),
            _ => Err(ErrorKind::UnknownStyleName(s.to_owned()).into()),
        }
    }
}

pub struct OutputComponents(HashSet<OutputComponent>);

impl OutputComponents {
    fn changes(&self) -> bool {
        self.0.contains(&OutputComponent::Changes)
    }

    fn grid(&self) -> bool {
        self.0.contains(&OutputComponent::Grid)
    }

    fn header(&self) -> bool {
        self.0.contains(&OutputComponent::Header)
    }

    fn numbers(&self) -> bool {
        self.0.contains(&OutputComponent::Numbers)
    }
}

pub struct Options<'a> {
    pub true_color: bool,
    pub output_components: OutputComponents,
    pub language: Option<&'a str>,
    pub colored_output: bool,
    pub paging: bool,
    pub term_width: usize,
}

enum OutputType {
    Pager(Child),
    Stdout(io::Stdout),
}

impl OutputType {
    fn pager() -> Result<Self> {
        Ok(OutputType::Pager(Command::new("less")
            .args(&["--quit-if-one-screen", "--RAW-CONTROL-CHARS", "--no-init"])
            .stdin(Stdio::piped())
            .spawn()
            .chain_err(|| "Could not spawn pager")?))
    }

    fn stdout() -> Self {
        OutputType::Stdout(io::stdout())
    }

    fn handle(&mut self) -> Result<&mut Write> {
        Ok(match *self {
            OutputType::Pager(ref mut command) => command
                .stdin
                .as_mut()
                .chain_err(|| "Could not open stdin for pager")?,
            OutputType::Stdout(ref mut handle) => handle,
        })
    }
}

impl Drop for OutputType {
    fn drop(&mut self) {
        if let OutputType::Pager(ref mut command) = *self {
            let _ = command.wait();
        }
    }
}

const GRID_COLOR: u8 = 238;
const LINE_NUMBER_COLOR: u8 = 244;

#[derive(Default)]
pub struct Colors {
    pub grid: Style,
    pub filename: Style,
    pub git_added: Style,
    pub git_removed: Style,
    pub git_modified: Style,
    pub line_number: Style,
}

impl Colors {
    fn plain() -> Self {
        Colors::default()
    }

    fn colored() -> Self {
        Colors {
            grid: Fixed(GRID_COLOR).normal(),
            filename: White.bold(),
            git_added: Green.normal(),
            git_removed: Red.normal(),
            git_modified: Yellow.normal(),
            line_number: Fixed(LINE_NUMBER_COLOR).normal(),
        }
    }
}

fn print_file(
    options: &Options,
    theme: &Theme,
    syntax_set: &SyntaxSet,
    printer: &mut Printer,
    filename: Option<&str>,
) -> Result<()> {
    let stdin = io::stdin(); // TODO: this is not always needed

    let mut reader: Box<BufRead> = match filename {
        None => Box::new(stdin.lock()),
        Some(filename) => Box::new(BufReader::new(File::open(filename)?)),
    };

    let syntax = match (options.language, filename) {
        (Some(language), _) => syntax_set.find_syntax_by_token(language),
        (None, Some(filename)) => {
            #[cfg(not(unix))]
            let may_read_from_file = true;

            // Do not peek at the file (to determine the syntax) if it is a FIFO because they can
            // only be read once.
            #[cfg(unix)]
            let may_read_from_file = !fs::metadata(filename)?.file_type().is_fifo();

            if may_read_from_file {
                syntax_set.find_syntax_for_file(filename)?
            } else {
                None
            }
        }
        (None, None) => None,
    };

    let syntax = syntax.unwrap_or_else(|| syntax_set.find_syntax_plain_text());
    let mut highlighter = HighlightLines::new(syntax, theme);

    printer.print_header(filename)?;

    let mut line_nr = 1;
    let mut line_buffer = String::new();
    loop {
        line_buffer.clear();
        let num_bytes = reader.read_line(&mut line_buffer);

        let line = match num_bytes {
            Ok(0) => {
                break;
            }
            Ok(_) => {
                if !line_buffer.ends_with("\n") {
                    line_buffer.push('\n');
                }
                &line_buffer
            }
            Err(_) => "<bat: INVALID UTF-8>\n",
        };

        let regions = highlighter.highlight(line);

        printer.print_line(line_nr, &regions)?;

        line_nr += 1;
    }

    printer.print_footer()?;

    Ok(())
}

fn get_output_type(paging: bool) -> OutputType {
    if paging {
        OutputType::pager().unwrap_or_else(|_| OutputType::stdout())
    } else {
        OutputType::stdout()
    }
}

fn is_truecolor_terminal() -> bool {
    env::var("COLORTERM")
        .map(|colorterm| colorterm == "truecolor" || colorterm == "24bit")
        .unwrap_or(false)
}

fn run() -> Result<()> {
    let interactive_terminal = atty::is(Stream::Stdout);

    let clap_color_setting = if interactive_terminal {
        AppSettings::ColoredHelp
    } else {
        AppSettings::ColorNever
    };

    let app_matches = App::new(crate_name!())
        .version(crate_version!())
        .global_setting(clap_color_setting)
        .global_setting(AppSettings::DeriveDisplayOrder)
        .global_setting(AppSettings::UnifiedHelpMessage)
        .global_setting(AppSettings::NextLineHelp)
        .setting(AppSettings::InferSubcommands)
        .setting(AppSettings::ArgsNegateSubcommands)
        .setting(AppSettings::DisableHelpSubcommand)
        .setting(AppSettings::VersionlessSubcommands)
        .max_term_width(90)
        .about(crate_description!())
        .arg(
            Arg::with_name("language")
                .short("l")
                .long("language")
                .help("Set the language for highlighting")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("FILE")
                .help("File(s) to print")
                .multiple(true)
                .empty_values(false),
        )
        .arg(
            Arg::with_name("style")
                .long("style")
                .use_delimiter(true)
                .takes_value(true)
                .possible_values(&["full", "plain", "changes", "header", "grid", "numbers"])
                .default_value("full")
                .help("Additional info to display along with content"),
        )
        .arg(
            Arg::with_name("color")
                .long("color")
                .takes_value(true)
                .possible_values(&["auto", "never", "always"])
                .default_value("auto")
                .help("When to use colors"),
        )
        .arg(
            Arg::with_name("paging")
                .long("paging")
                .takes_value(true)
                .possible_values(&["auto", "never", "always"])
                .default_value("auto")
                .help("When to use the pager"),
        )
        .arg(
            Arg::with_name("list-languages")
                .long("list-languages")
                .help("Displays supported languages"),
        )
        .subcommand(
            SubCommand::with_name("cache")
                .about("Modify the syntax-definition and theme cache")
                .arg(
                    Arg::with_name("init")
                        .long("init")
                        .short("i")
                        .help("Initialize the cache by loading from the config dir"),
                )
                .arg(
                    Arg::with_name("clear")
                        .long("clear")
                        .short("c")
                        .help("Reset the cache"),
                )
                .arg(
                    Arg::with_name("config-dir")
                        .long("config-dir")
                        .short("d")
                        .help("Show the configuration directory"),
                )
                .group(
                    ArgGroup::with_name("cache-actions")
                        .args(&["init", "clear", "config-dir"])
                        .required(true),
                ),
        )
        .help_message("Print this help message.")
        .version_message("Show version information.")
        .get_matches();

    match app_matches.subcommand() {
        ("cache", Some(cache_matches)) => {
            if cache_matches.is_present("init") {
                let assets = HighlightingAssets::from_files()?;
                assets.save()?;
            } else if cache_matches.is_present("clear") {
                print!("Clearing theme set cache ... ");
                fs::remove_file(theme_set_path())?;
                println!("okay");

                print!("Clearing syntax set cache ... ");
                fs::remove_file(syntax_set_path())?;
                println!("okay");
            } else if cache_matches.is_present("config-dir") {
                println!("{}", config_dir());
            }
        }
        _ => {
            let files: Vec<Option<&str>> = app_matches
                .values_of("FILE")
                .map(|values| {
                    values
                        .map(|filename| {
                            if filename == "-" {
                                None
                            } else {
                                Some(filename)
                            }
                        })
                        .collect()
                })
                .unwrap_or_else(|| vec![None]); // read from stdin (None) if no args are given

            let output_components = values_t!(app_matches.values_of("style"), OutputComponent)?
                .into_iter()
                .map(|style| style.components())
                .fold(HashSet::new(), |mut acc, components| {
                    acc.extend(components.iter().cloned());
                    acc
                });

            let options = Options {
                true_color: is_truecolor_terminal(),
                output_components: OutputComponents(output_components),
                language: app_matches.value_of("language"),
                colored_output: match app_matches.value_of("color") {
                    Some("always") => true,
                    Some("never") => false,
                    _ => interactive_terminal,
                },
                paging: match app_matches.value_of("paging") {
                    Some("always") => true,
                    Some("never") => false,
                    Some("auto") | _ => if files.contains(&None) {
                        // If we are reading from stdin, only enable paging if we write to an
                        // interactive terminal and if we do not *read* from an interactive
                        // terminal.
                        interactive_terminal && !atty::is(Stream::Stdin)
                    } else {
                        interactive_terminal
                    },
                },
                term_width: console::Term::stdout().size().1 as usize,
            };

            let assets = HighlightingAssets::new();
            let theme = assets.default_theme()?;

            if app_matches.is_present("list-languages") {
                let languages = assets.syntax_set.syntaxes();

                let longest = languages
                    .iter()
                    .filter(|s| !s.hidden && s.file_extensions.len() > 0)
                    .map(|s| s.name.len())
                    .max()
                    .unwrap_or(32); // Fallback width if they have no language definitions.

                let separator = " ";
                for lang in languages {
                    if lang.hidden || lang.file_extensions.len() == 0 {
                        continue;
                    }
                    print!("{:width$}{}", lang.name, separator, width = longest);

                    // Line-wrapping for the possible file extension overflow.
                    let desired_width = options.term_width - longest - separator.len();
                    // Number of characters on this line so far, wrap before `desired_width`
                    let mut num_chars = 0;

                    let comma_separator = ", ";
                    let mut extension = lang.file_extensions.iter().peekable();
                    while let Some(word) = extension.next() {
                        // If we can't fit this word in, then create a line break and align it in.
                        let new_chars = word.len() + comma_separator.len();
                        if num_chars + new_chars >= desired_width {
                            num_chars = 0;
                            print!("\n{:width$}{}", "", separator, width = longest);
                        }

                        num_chars += new_chars;
                        print!("{}", Green.paint(word as &str));
                        if extension.peek().is_some() {
                            print!("{}", comma_separator);
                        }
                    }
                    println!();
                }

                return Ok(());
            }

            let mut output_type = get_output_type(options.paging);
            let handle = output_type.handle()?;
            let mut printer = Printer::new(handle, &options);

            for file in files {
                printer.line_changes = file.and_then(|filename| get_git_diff(filename));

                print_file(&options, theme, &assets.syntax_set, &mut printer, file)?;
            }
        }
    }

    Ok(())
}

fn main() {
    let result = run();

    if let Err(error) = result {
        match error {
            Error(ErrorKind::Io(ref io_error), _)
                if io_error.kind() == io::ErrorKind::BrokenPipe => {}
            _ => {
                eprintln!("{}: {}", Red.paint("[bat error]"), error);

                process::exit(1);
            }
        };
    }
}