chop/src/main.rs

use std::error::Error;
use std::fs::File;
use std::io::{self, BufReader};
use std::io::prelude::*;

use argh::FromArgs;

#[derive(FromArgs)]
#[argh(
    description = "\
    chop off bits of file\n\n\
    chop will take produce --count files of --lines lines from the beginning of a\n\
    file or stdin. Any remaining lines will be written to a final catchall file.")]
struct Options {

    /// number of lines in each chunk
    #[argh(option, short = 'n')]
    lines: usize,

    /// count of chunks to produce (default 1)
    #[argh(option, short = 'c', default = "1")]
    count: usize,

    /// optional prefix to use for generated files (default "")
    #[argh(option, short = 'p', default = "String::from(\"\")")]
    prefix: String,

    /// use numeric suffixes starting with 0, not alphabetic 
    #[argh(switch, short = 'd')]
    numeric: bool,

    /// use numeric suffixes starting with 0, not alphabetic 
    #[argh(option, long = "numeric-start", default = "0")]
    numeric_start: usize,

    /// filename to read from, or "-" for stdin (default "-")
    #[argh(positional, default = "String::from(\"-\")")]
    filename: String,


}

struct AlphabeticSuffixGenerator {
    suffix: Vec<u8>,
    remaining: usize,
}

impl AlphabeticSuffixGenerator {
    fn new(count: usize) -> AlphabeticSuffixGenerator {
        // given 27 letters, minimum number of digits is floor( Log 27 (count) + 1 )
        let suffix_length =  ((count as f64).log(27f64) + 1.0).floor() as usize;

        let mut suffix = vec![b'a'; suffix_length];
        suffix[suffix_length-1] -= 1;

        AlphabeticSuffixGenerator{suffix, remaining: count}
    }
}

impl Iterator for AlphabeticSuffixGenerator {
    type Item = String;

    fn next(&mut self) -> Option<String> {
        if self.remaining == 0 {
            return None
        }
        for idx in (0..self.suffix.len()).rev() {
            if self.suffix[idx] < b'z' {
                self.suffix[idx] += 1;
                break
            } else {
                self.suffix[idx] = b'a';
            }
        }
        self.remaining -= 1;
        Some(String::from_utf8(self.suffix.to_vec()).expect("invalid suffix generated"))
    }
}


fn try_main() -> Result<(), Box<dyn Error>> {
    let opts: Options = argh::from_env();

    let input: Box<dyn Read> = match opts.filename.as_str() {
        "-" => Box::new(io::stdin()),
         path => { 
             match File::open(path) {
                 Err(why) => panic!("failed to open {}: {}", path, why),
                 Ok(file) => Box::new(file),
            }
         },
    };

    let prefix = match opts.prefix.as_str() {
        "" => String::from(""),
        _ => opts.prefix + "_",
    };
   
    let suffix_gen: Box<dyn Iterator<Item = String>> = if opts.numeric {
        Box::new((opts.numeric_start..(opts.numeric_start + opts.count)).map(|num| format!("{}", num)))
    } else {
        Box::new(AlphabeticSuffixGenerator::new(opts.count))
    };

    let mut eof = false;
    let mut reader = BufReader::new(input);
    for suffix in suffix_gen {
        let mut out_file = File::create(format!("{}{}", prefix, suffix))?;
        let mut line = String::new();
        for _ in 0..opts.lines {
            match reader.read_line(&mut line)? {
                0 => {
                    eof = true;
                    break
                },
                _ => {
                    out_file.write_all(line.as_bytes())?;
                    line.clear();
                }
            };
        }
    }
    // see if there's anything left for the remainder file
    if !eof {
        let mut out_file = File::create(format!("{}rest", prefix))?;
        for result in reader.lines() {
            // lines() strips newline characters, so add one
            out_file.write_all((result?+"\n").as_bytes())?;
        };
    }
    
    Ok(())
}

fn main() {
    // display nicer, non-debug representation of the error
    if let Err(err) = try_main() {
        eprintln!("{}", err);
        std::process::exit(1);
    }
}