1use crate::errors::{FormatError, SourceError};
2use crate::input::{Formatter, SortOrder, Source};
3use anyhow::Result;
4use regex::Regex;
5use std::fs;
6use std::iter::{IntoIterator, Iterator};
7use std::path::Path;
8use std::vec::IntoIter;
9use walkdir::IntoIter as WalkIter;
10use walkdir::WalkDir;
11
12pub enum InputIterator {
13 VectorIterator(IntoIter<(String, String)>),
14 DirectoryIterator {
15 formatter: Formatter,
16 re: Regex,
17 preserve_extension: bool,
18 iter: WalkIter,
19 },
20}
21
22impl InputIterator {
23 pub fn new(
24 source: Source,
25 formatter: Option<Formatter>,
26 preserve_extension: bool,
27 ) -> Result<Self> {
28 if let Source::Map(map) = source {
29 return Ok(Self::VectorIterator(map.into_iter()));
30 }
31
32 let formatter = formatter.ok_or(FormatError::EmptyFormatter)?;
33
34 if let Source::Sort(order) = source {
35 let mut map = Vec::new();
36 let mut inputs = Vec::new();
37 for entry in fs::read_dir(".")?.flatten() {
38 inputs.push(entry.file_name().to_string_lossy().to_string());
39 }
40 inputs.sort_by(|a, b| {
41 if order == SortOrder::Asc {
42 natord::compare(a, b)
43 } else {
44 natord::compare(b, a)
45 }
46 });
47 for (i, input) in inputs.into_iter().enumerate() {
48 let index = (i + 1).to_string();
49 let mut output = formatter.format(vec![input.as_str(), index.as_str()].as_slice());
50 if preserve_extension {
51 if let Some(extension) = Path::new(input.as_str()).extension() {
52 output.push('.');
53 output.push_str(extension.to_str().unwrap_or_default());
54 }
55 }
56 map.push((input, output));
57 }
58 return Ok(Self::VectorIterator(map.into_iter()));
59 }
60
61 if let Source::Regex(re, depth, max_depth) = source {
62 let max_depth = max_depth.unwrap_or(depth);
63 return Ok(Self::DirectoryIterator {
64 formatter,
65 re,
66 preserve_extension,
67 iter: WalkDir::new(".")
68 .min_depth(if depth > max_depth { max_depth } else { depth })
69 .max_depth(max_depth)
70 .into_iter(),
71 });
72 }
73
74 Err(SourceError::new(String::from("unknown source")).into())
75 }
76}
77
78impl Iterator for InputIterator {
79 type Item = (String, String);
80
81 fn next(&mut self) -> Option<Self::Item> {
82 match self {
83 Self::VectorIterator(iter) => iter.next(),
84 Self::DirectoryIterator {
85 formatter,
86 re,
87 preserve_extension,
88 iter,
89 } => {
90 for entry in iter {
91 let Ok(entry) = entry else {
92 continue;
93 };
94 let path = entry.path();
95 let input = path.strip_prefix("./").unwrap_or(path).to_string_lossy();
96 let Some(cap) = re.captures(input.as_ref()) else {
97 continue;
98 };
99 let vars: Vec<&str> = cap
100 .iter()
101 .map(|c| c.map(|c| c.as_str()).unwrap_or_default())
102 .collect();
103 let mut output = formatter.format(vars.as_slice());
104 if *preserve_extension {
105 if let Some(extension) = Path::new(input.as_ref()).extension() {
106 output.push('.');
107 output.push_str(extension.to_str().unwrap_or_default());
108 }
109 }
110 return Some((input.to_string(), output));
111 }
112 None
113 }
114 }
115 }
116}