shitbox/corebox/commands/head/mod.rs

use super::Cmd;
use clap::{value_parser, Arg, ArgAction, ArgMatches, Command};
use shitbox::{args, Path};
use std::{
    env,
    error::Error,
    fs,
    io::{stdin, Read, Write},
};
use termcolor::{Color, ColorChoice, ColorSpec, StandardStream, WriteColor};

#[derive(Debug)]
pub struct Head;

impl Cmd for Head {
    fn cli(&self) -> Command {
        Command::new("head")
            .author("Nathan Fisher")
            .about("Display first lines of a file")
            .long_about(
                "Print the first 10 lines of each FILE to standard output.\n\
                With more than one FILE, precede each with a header giving the file name.\n\n\
                With no FILE, or when FILE is -, read standard input.",
            )
            .args([
                Arg::new("FILES")
                    .help("The input file to use")
                    .default_value("-")
                    .num_args(1..),
                Arg::new("BYTES")
                    .help("Count bytes instead of lines")
                    .short('c')
                    .long("bytes")
                    .action(ArgAction::SetTrue),
                Arg::new("QUIET")
                    .help("Disable printing a header. Overrides -c")
                    .short('q')
                    .long("quiet")
                    .action(ArgAction::SetTrue),
                args::header(),
                Arg::new("LINES")
                    .help("Count n number of lines (or bytes if -c is specified).")
                    .short('n')
                    .long("lines")
                    .allow_negative_numbers(false)
                    .conflicts_with("num")
                    .value_parser(value_parser!(usize)),
                Arg::new("color")
                    .short('C')
                    .long("color")
                    .value_parser(["always", "ansi", "auto", "never"]),
                Arg::new("num")
                    .short('1')
                    .short_aliases(['2', '3', '4', '5', '6', '7', '8', '9'])
                    .hide(true)
                    .action(ArgAction::Append),
            ])
    }

    fn run(&self, matches: &ArgMatches) -> Result<(), Box<dyn Error>> {
        let args: Vec<_> = env::args().collect();
        let idx = match shitbox::progname() {
            Some(s) if s.as_str() == "head" => 1,
            _ => 2,
        };
        let mut lines = 10;
        if args.len() > idx {
            let one = &args[idx];
            if one.starts_with('-') && one.len() > 1 {
                if let Ok(lines) = one[1..].parse::<usize>() {
                    let files: Vec<_> = if args.len() > idx + 1 {
                        args[idx + 1..].iter().map(String::as_str).collect()
                    } else {
                        vec!["-"]
                    };
                    for file in &files {
                        head(file, lines, false, false, ColorChoice::Never)?;
                    }
                    return Ok(());
                }
            }
        }
        if let Some(l) = matches.get_one("LINES") {
            lines = *l;
        }
        let color = match matches.get_one::<String>("color").map(String::as_str) {
            Some("always") => ColorChoice::Always,
            Some("ansi") => ColorChoice::AlwaysAnsi,
            Some("auto") => {
                if atty::is(atty::Stream::Stdout) {
                    ColorChoice::Auto
                } else {
                    ColorChoice::Never
                }
            }
            _ => ColorChoice::Never,
        };
        if let Some(files) = matches.get_many::<String>("FILES") {
            let (len, _) = files.size_hint();
            let header = !matches.get_flag("QUIET") && { len > 1 || matches.get_flag("HEADER") };
            for (index, file) in files.enumerate() {
                if index == 1 && header {
                    println!();
                }
                head(file, lines, header, matches.get_flag("BYTES"), color)?;
            }
        }
        Ok(())
    }

    fn path(&self) -> Option<Path> {
        Some(shitbox::Path::Bin)
    }
}

fn head(
    file: &str,
    count: usize,
    header: bool,
    bytes: bool,
    color: ColorChoice,
) -> Result<(), Box<dyn Error>> {
    let mut contents = String::new();
    if file == "-" {
        stdin().read_to_string(&mut contents)?;
    } else {
        contents = fs::read_to_string(file)?;
    }
    if header {
        let mut stdout = StandardStream::stdout(color);
        stdout.set_color(ColorSpec::new().set_fg(Some(Color::Green)))?;
        writeln!(stdout, "==> {file} <==")?;
        stdout.reset()?;
    }
    if bytes {
        for (index, char) in contents.chars().enumerate() {
            if index < count {
                print!("{char}");
            } else {
                println!();
                return Ok(());
            }
        }
        println!();
    } else {
        for (index, line) in contents.lines().enumerate() {
            if index < count {
                println!("{line}");
            } else {
                return Ok(());
            }
        }
    }
    Ok(())
}