simino/input/
iterator.rs

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}