shitbox/src/cmd/fold/mod.rs

use super::Cmd;
use clap::{value_parser, Arg, ArgAction, ArgMatches, Command};
use std::{fs::File, io::{self, BufRead, BufReader}};
use textwrap::{fill, wrap_algorithms::WrapAlgorithm};

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

impl Default for Fold {
    fn default() -> Self {
        Self {
            name: "fold",
            path: Some(crate::Path::UsrBin),
        }
    }
}

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

    fn cli(&self) -> Command {
        Command::new(self.name)
            .about("Wrap each input line to fit in specified width")
            .author("Nathan Fisher")
            .after_long_help("With no FILE, or when FILE is -, read standard input")
            .args([
                Arg::new("FILE").help("The input file to use").num_args(0..),
                Arg::new("BYTES")
                    .help("Count bytes rather than columns")
                    .short('b')
                    .long("bytes")
                    .action(ArgAction::SetTrue),
                Arg::new("WORDS")
                    .help("Break at spaces")
                    .short('s')
                    .long("spaces")
                    .action(ArgAction::SetTrue),
                Arg::new("OPTIMAL")
                    .help("Optimal fit")
                    .long_help(
                        "Uses a look ahad algorithm to avoid unnecessarily short or long lines",
                    )
                    .short('o')
                    .long("optimal")
                    .action(ArgAction::SetTrue)
                    .conflicts_with_all(["BYTES", "WORDS"]),
                Arg::new("WIDTH")
                    .help("Use width columns")
                    .short('w')
                    .long("width")
                    .default_value("80")
                    .num_args(1)
                    .value_parser(value_parser!(usize)),
            ])
    }

    fn run(&self, matches: Option<&ArgMatches>) -> Result<(), Box<dyn std::error::Error>> {
        let Some(matches) = matches else {
            return Err(Box::new(io::Error::new(io::ErrorKind::Other, "no input")));
        };
        let files = match matches.get_many::<String>("FILE") {
            Some(c) => c.map(String::to_string).collect(),
            None => vec!["-".to_string()],
        };
        for file in files {
            if file.as_str() == "-" {
                wrap_stdin(matches)?;
            } else {
                wrap_file(&file, matches)?;
            }
        }
        if matches.get_flag("BYTES") {
            println!();
        }
        Ok(())
    }

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

fn wrap_line(line: &str, args: &ArgMatches) {
    let width = args.get_one("WIDTH").map(|x| *x).unwrap_or(80);
    if args.get_flag("OPTIMAL") {
        let line = line.replace('\t', "    ");
        let opts = textwrap::Options::new(width).wrap_algorithm(WrapAlgorithm::new_optimal_fit());
        println!("{}", fill(line.trim_end(), opts));
    } else if args.get_flag("WORDS") {
        let line = line.replace('\t', "    ");
        let opts = textwrap::Options::new(width).wrap_algorithm(WrapAlgorithm::FirstFit);
        println!("{}", fill(line.trim_end(), opts));
    } else if args.get_flag("BYTES") {
        for (index, b) in line.as_bytes().iter().enumerate() {
            if index % width == 0 {
                println!();
            }
            print!("{}", *b as char);
        }
    } else {
        let line = line
            .chars()
            .collect::<Vec<char>>()
            .chunks(width)
            .map(|c| c.iter().collect::<String>())
            .collect::<Vec<String>>();
        for line in &line {
            println!("{}", line);
        }
    }
}

fn wrap_stdin(args: &ArgMatches) -> Result<(), io::Error> {
    for line in io::stdin().lock().lines() {
        wrap_line(&line?, args);
    }
    Ok(())
}

fn wrap_file(file: &str, args: &ArgMatches) -> Result<(), io::Error> {
    let fd = File::open(file)?;
    let buf = BufReader::new(fd);
    for line in buf.lines() {
        wrap_line(&line?, args);
    }
    Ok(())
}