diff --git a/src/bin/counter.rs b/src/bin/counter.rs index 05a314b..92bff17 100644 --- a/src/bin/counter.rs +++ b/src/bin/counter.rs @@ -1,22 +1,42 @@ -use clap::Parser; -use std::path::PathBuf; +use clap::{CommandFactory, Parser}; +use std::{ + fs::File, + io::{stdin, BufRead, BufReader, IsTerminal}, + path::PathBuf, +}; /// Count the number of lines in a file #[derive(Parser)] #[command(arg_required_else_help = true)] struct Cli { - /// The path to the file to read + /// The path to the file to read, use - to read from stdin (must not be a tty) file: PathBuf, } fn main() { let args = Cli::parse(); - let mut word_count = 0; - let file = args.file; + let word_count; + let mut file = args.file; - for line in std::fs::read_to_string(&file).unwrap().lines() { - word_count += line.split(' ').count(); + if file == PathBuf::from("-") { + if stdin().is_terminal() { + Cli::command().print_help().unwrap(); + ::std::process::exit(2); + } + + file = PathBuf::from(""); + word_count = words_in_buf_reader(BufReader::new(stdin().lock())); + } else { + word_count = words_in_buf_reader(BufReader::new(File::open(&file).unwrap())); } - println!("Words in {}: {}", file.to_str().unwrap(), word_count) + println!("Words in {}: {}", file.to_string_lossy(), word_count) +} + +fn words_in_buf_reader(buf_reader: R) -> usize { + let mut count = 0; + for line in buf_reader.lines() { + count += line.unwrap().split(' ').count() + } + count }