shitbox/src/cmd/head/mod.rs

use super::Cmd;
use crate::Path;
use clap::{value_parser, Arg, ArgAction, ArgMatches, Command};
use std::{
    error::Error,
    fs,
    io::{self, stdin, Read},
    process,
};

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)
        .version(env!("CARGO_PKG_VERSION"))
        .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\
            With no FILE, or when FILE is -, read standard input.\n\
            Mandatory arguments to long options are mandatory for short options too."
        )
        .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")
                .default_value("10")
                .allow_negative_numbers(false)
                .value_parser(value_parser!(usize))
                ])
    }

    fn run(&self, matches: Option<&ArgMatches>) -> Result<(), Box<dyn Error>> {
        let matches = if let Some(m) = matches {
            m
        } else {
            return Err(io::Error::new(io::ErrorKind::Other, "No input").into());
        };
        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,
                match matches.get_one("LINES") {
                    Some(c) => *c,
                    None => 10,
                },
                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;
            }
        }
    }
}