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},
    process,
};

#[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("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 {
                let lines: usize = arg1[1..].parse()?;
                let files: Vec<_> = if args.len() > idx + 1 {
                    args[idx + 1..].iter().map(|x| x.as_str()).collect()
                } else {
                    vec!["-"]
                };
                for file in &files {
                    head(&file, lines, false, false);
                }
                return Ok(());
            }
        }
        let matches = if let Some(m) = matches {
            m
        } 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") };
        for (index, file) in files.into_iter().enumerate() {
            if index == 1 && header {
                println!();
            }
            head(&file, lines, header, matches.get_flag("BYTES"));
        }
        Ok(())
    }

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

fn head(file: &str, count: usize, header: bool, bytes: bool) {
    let mut contents = String::new();
    if file == "-" {
        match stdin().read_to_string(&mut contents) {
            Ok(_) => true,
            Err(e) => {
                eprintln!("head: {e}");
                process::exit(1);
            }
        };
    } else {
        let buf = fs::read_to_string(file);
        contents = match buf {
            Ok(c) => c,
            Err(e) => {
                eprintln!("head: {e}");
                process::exit(1);
            }
        };
    }
    if header {
        println!("==> {file} <==");
    }
    if bytes {
        for (index, char) in contents.chars().into_iter().enumerate() {
            if index < count {
                print!("{char}");
            } else {
                println!();
                return;
            }
        }
        println!();
    } else {
        for (index, line) in contents.lines().into_iter().enumerate() {
            if index < count {
                println!("{line}");
            } else {
                return;
            }
        }
    }
}