shitbox/src/cmd/head/mod.rs

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

#[derive(Debug)]
pub struct Head {
    name: &'static str,
    path: Option<Path>,
}

pub const HEAD: Head = Head {
    name: "head",
    path: Some(Path::Bin),
};

impl Cmd for Head {
    fn name(&self) -> &str {
        self.name
    }

    fn cli(&self) -> Command {
        Command::new(self.name)
            .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")
                    .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),
        		Arg::new("HEADER")
        		    .help("Each file is preceded by a header consisting of the string \"==> XXX <==\" where \"XXX\" is the name of the file.")
    	    	    .short('v')
    		        .long("verbose")
                    .action(ArgAction::SetTrue),
                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: Option<&ArgMatches>) -> Result<(), Box<dyn Error>> {
        let args: Vec<_> = env::args().collect();
        let idx = match crate::progname() {
            Some(s) if s.as_str() == "head" => 1,
            _ => 2,
        };
        let mut lines = 10;
        if args.len() > idx {
            let arg1 = &args[idx];
            if arg1.starts_with('-') && arg1.len() > 1 {
                if let Ok(lines) = arg1[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(());
                }
            }
        }
        let Some(matches) = matches else {
            return Err(io::Error::new(io::ErrorKind::Other, "No input").into());
        };
        if let Some(l) = matches.get_one("LINES") {
            lines = *l;
        }
        let files = match matches.get_many::<String>("FILES") {
            Some(c) => c.map(std::string::ToString::to_string).collect(),
            None => vec!["-".to_string()],
        };
        let header =
            !matches.get_flag("QUIET") && { files.len() > 1 || matches.get_flag("HEADER") };
        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,
        };
        for (index, file) in files.into_iter().enumerate() {
            if index == 1 && header {
                println!();
            }
            head(&file, lines, header, matches.get_flag("BYTES"), color)?;
        }
        Ok(())
    }

    fn path(&self) -> Option<Path> {
        self.path
    }
}

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(())
}