--only and count blank lines

src/fc.rs

@@ -10,6 +10,7 @@ pub struct FileContent {
 	pub path: PathBuf,
 	pub language: Language,
 	pub lines: usize,
+	pub blank_lines: usize,
 }
 
 impl FileContent {
@@ -23,12 +24,21 @@ impl FileContent {
 		))?;
 
 		let text = read_to_string(&path)?;
-		let lines = text.lines().count();
+		let mut lines = 0;
+		let mut blank_lines = 0;
+
+		for line in text.lines() {
+			if line.is_empty() {
+				blank_lines += 1;
+			}
+			lines += 0;
+		}
 
 		Ok(Self {
 			path,
 			language,
 			lines,
+			blank_lines,
 		})
 	}
 }

src/langs.rs

@@ -577,6 +577,7 @@ impl LanguageInfo {
 pub struct LanguageSummary {
 	pub language: Language,
 	pub lines: usize,
+	pub blank_lines: usize,
 	pub files: Vec<PathBuf>,
 }
 
@@ -585,6 +586,7 @@ impl LanguageSummary {
 		Self {
 			language,
 			lines: 0,
+			blank_lines: 0,
 			files: vec![],
 		}
 	}

src/options.rs

@@ -20,6 +20,7 @@ pub struct Options {
 	pub blame: bool,
 	pub head: Option<usize>,
 	pub excluded: HashSet<Language>,
+	pub only_include: HashSet<Language>,
 	pub total_lines_only: bool,
 }
 
@@ -40,6 +41,7 @@ impl Default for Options {
 			blame: false,
 			head: None,
 			excluded: Default::default(),
+			only_include: Default::default(),
 			total_lines_only: false,
 		}
 	}
@@ -125,6 +127,21 @@ where
 						);
 					}
 				}
+				"-o" | "-only" | "--only" => {
+					let include = args.next();
+					let list = include
+						.as_ref()
+						.expect(&format!("expected a language to follow {} flag", arg))
+						.as_ref()
+						.split(",");
+					for lang in list {
+						options.only_include.insert(
+							Language::from_name(lang)
+								.or_else(|| Language::from_extension(OsStr::new(lang)))
+								.expect(&format!("unrecognized language identifier \"{}\"", lang)),
+						);
+					}
+				}
 				"-l" | "-lines" | "--lines" => {
 					if options.head.is_some() {
 						println!("{} is incompatible with -t/--top", arg);
@@ -140,6 +157,10 @@ where
 			}
 		}
 
+		if !options.only_include.is_empty() && !options.excluded.is_empty() {
+			println!("warning: both --only and --exclude have been set, which doesn't really make sense")
+		}
+
 		options
 	}
 }

src/scan.rs

@@ -12,7 +12,7 @@ use crate::langs::LanguageSummary;
 use crate::options::Options;
 
 pub fn scan(options: Options) -> anyhow::Result<()> {
-	let mut summary: HashMap<Language, LanguageSummary> = Default::default();
+	let mut summaries: HashMap<Language, LanguageSummary> = Default::default();
 	let dir = &OsString::from(&options.root_dir);
 	let dir_path = Path::new(dir);
 
@@ -42,37 +42,35 @@ pub fn scan(options: Options) -> anyhow::Result<()> {
 			continue;
 		};
 
-		let summary = summary
+		let summary = summaries
 			.entry(content.language)
 			.or_insert_with(|| LanguageSummary::from(content.language));
 		summary.lines += content.lines;
 		summary.files.push(path);
 	}
 
-	let mut summary = summary.iter().collect::<Vec<_>>();
-	summary.sort_by(|a, b| b.1.lines.cmp(&a.1.lines));
+	let mut summaries = summaries.iter().collect::<Vec<_>>();
+	summaries.sort_by(|a, b| b.1.lines.cmp(&a.1.lines));
 
-	let result_iter = || {
-		let mut count = 0;
-		summary
-			.iter()
-			.filter(|(lang, _)| !options.excluded.contains(lang))
-			.take_while(move |_| {
-				if let Some(max) = &options.head {
-					if count >= *max {
-						return false;
-					}
-
-					count += 1;
-				}
-				true
-			})
-	};
+	if !options.excluded.is_empty() {
+		summaries.retain(|(lang, _)| !options.excluded.contains(lang))
+	}
+
+	if !options.only_include.is_empty() {
+		summaries.retain(|(lang, _)| options.only_include.contains(lang))
+	}
+
+	if let Some(max) = &options.head {
+		summaries.truncate(*max);
+	}
 
 	// We wait until here to handle `--lines` because it allows us to still respect
 	// other options like `--exclude` and `--top`.
 	if options.total_lines_only {
-		let total_lines = result_iter().map(|summary| summary.1.lines).sum::<usize>();
+		let total_lines = summaries
+			.iter()
+			.map(|(_, summary)| summary.lines)
+			.sum::<usize>();
 
 		println!("{}", total_lines);
 		return Ok(());
@@ -81,28 +79,29 @@ pub fn scan(options: Options) -> anyhow::Result<()> {
 	let inner_width = options.width - 2; // we have a padding of 1 character on each side
 
 	println!();
-	result_iter().for_each(|(_, stat)| {
+	for (_, summary) in summaries.iter() {
 		println!(
 			" {:width$}",
-			stat.to_terminal_display(&options),
+			summary.to_terminal_display(&options),
 			width = inner_width
 		)
-	});
+	}
 
-	let total_lines = result_iter()
-		.map(|(_, stat)| stat.lines)
+	let total_lines = summaries
+		.iter()
+		.map(|(_, summary)| summary.lines)
 		.reduce(|acc, lines| acc + lines)
 		.ok_or_else(|| anyhow!("no code found in \"{}\"", dir_path.display()))?;
 
 	let mut filled = 0;
 
-	result_iter().for_each(|(_, stat)| {
+	for (_, stat) in summaries.iter() {
 		// If there are 0 total lines, then just say everything is 0%.
 		let percent = (stat.lines * inner_width)
 			.checked_div(total_lines)
 			.unwrap_or(0);
 		if percent == 0 {
-			return;
+			continue;
 		}
 
 		// Print padding and such on first fill
@@ -117,7 +116,7 @@ pub fn scan(options: Options) -> anyhow::Result<()> {
 			Some(color) => print!("{}", color.on_color(&*" ".repeat(percent))),
 			None => print!("{}", " ".repeat(percent).on_white()),
 		};
-	});
+	}
 
 	// Don't print a bar at all if it'd just all be uncategorized.
 	if filled != 0 {

tests/main.rs

@@ -61,6 +61,19 @@ fn self_check_exclude() {
 	assert!(!stdout.contains("TOML"));
 }
 
+#[test]
+fn self_check_only() {
+	setup::before();
+
+	let result = Command::new(EXE).args(["-o", "rs"]).output().unwrap();
+	assert!(result.status.success());
+	let stdout = String::from_utf8_lossy(&result.stdout);
+
+	assert!(stdout.contains("Rust"));
+	assert!(!stdout.contains("Markdown"));
+	assert!(!stdout.contains("TOML"));
+}
+
 #[test]
 fn scan_nonexistent() {
 	setup::before();