1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70

use anyhow::Error;
use std::{env::current_dir, ffi::OsStr, fs, path::PathBuf};

use itertools::Either;
use std::collections::HashSet;
use walkdir::WalkDir;

use crate::constant::FILE_EXTENSION;

#[derive(Clone, Debug)]
pub struct SourceFiles {
    base: PathBuf,
    files: HashSet<PathBuf>,
}

impl SourceFiles {
    pub fn new(path: Option<PathBuf>) -> Result<SourceFiles, Error> {
        let base = match path {
            Some(path) => {
                if path.is_file() {
                    path.parent().unwrap().to_path_buf()
                } else if path.is_dir() {
                    path
                } else {
                    return Err(Error::msg("PathNotFound"));
                }
            }
            None => current_dir().unwrap(),
        };

        let file_paths = fs::read_dir(&base)
            .unwrap()
            .into_iter()
            .flat_map(|item| {
                let item = item.unwrap();

                if item.path().is_dir() {
                    Either::Left(
                        WalkDir::new(item.path())
                            .into_iter()
                            .filter_map(|e| e.ok())
                            .filter(move |e| {
                                e.path().extension() == Some(OsStr::new(FILE_EXTENSION))
                            })
                            .map(|e| e.into_path()),
                    )
                } else {
                    Either::Right(Box::new(
                        vec![item.path()]
                            .into_iter()
                            .filter(|e| e.extension() == Some(OsStr::new(FILE_EXTENSION))),
                    ))
                }
            })
            .collect::<HashSet<PathBuf>>();

        Ok(SourceFiles {
            base,
            files: file_paths,
        })
    }

    pub fn files(&self) -> &HashSet<PathBuf> {
        &self.files
    }

    pub fn base(&self) -> &PathBuf {
        &self.base
    }
}