tre/
streamer.rs

1use config::Options;
2use display;
3use ignore::{DirEntry, Error, Walk};
4use stats_collector::{FileType, StatsCollector};
5use walker::build_shallow;
6
7/// Streamer represents the object traversing the filesystem, printing the structure and collecting stats.
8pub struct Streamer {
9    prev_depth: usize,
10    curr_depth: usize,
11    parent_depths: Vec<usize>,
12    curr_line_count: Option<usize>,
13    collector: StatsCollector,
14    dir_only: bool,
15    count_lines: bool,
16    colours: bool,
17    options: Options,
18}
19
20impl Streamer {
21    pub fn new(opt: Options, collector: StatsCollector) -> Streamer {
22        Streamer {
23            prev_depth: 0,
24            curr_depth: 0,
25            parent_depths: vec![],
26            curr_line_count: None,
27            collector: collector,
28            dir_only: opt.dir_only,
29            count_lines: opt.line_count,
30            colours: !opt.no_colours,
31            options: opt,
32        }
33    }
34
35    /// kicks off a recursive streaming of a directory file structure
36    pub fn stream_tree(&mut self, walker: &mut Walk) -> Result<(), Error> {
37        let mut prev = walker
38            .next()
39            .expect("could not get first element from walker")
40            .expect("could not get first element from walker");
41        self.prev_depth = prev.depth();
42
43        // walker traverses depth first, ignoring files/folders it can't parse or does not have permission to
44        for dir in walker.into_iter().filter_map(|e| e.ok()) {
45            self.curr_depth = dir.depth();
46            self.stream_node(&prev, false)?;
47            prev = dir;
48            self.prev_depth = self.curr_depth;
49        }
50
51        self.stream_node(&prev, true)?;
52
53        println!("{}", self.collector);
54        Ok(())
55    }
56
57    /// parse and stream an individual node, correctly printing its representation and updating statistics.
58    fn stream_node(&mut self, node: &DirEntry, is_last: bool) -> Result<(), Error> {
59        let mut is_last = is_last;
60        let file_name = node.file_name().to_owned().into_string().unwrap();
61
62        //parses current file type and tally stats
63        let file_type = self.collector.parse_and_collect(node)?;
64
65        // match on file type to do additional logic, such as skip directory or do line counting
66        match file_type {
67            FileType::File if self.dir_only => return Ok(()),
68            FileType::File if self.count_lines => {
69                // if failed to count lines (.e.g. no permission to read file) just pretend its 0
70                self.curr_line_count = Some(self.collector.count_lines(node).unwrap_or(0));
71            }
72            _ => self.curr_line_count = None,
73        }
74
75        let mut should_pop = false;
76        // This logic allows us to keep record parents (store in vec) of the current file.
77        // We are always traversing one file ahead of what we print, so we can tell whether the thing to print
78        // is the last of its parent directory (when the depth changes)
79        // additional we build a shallow walker that tries to figure out if the paren dir is the last dir,
80        // if so we pop it from the vec stack because we don't need to print that branch afterwards.
81        if self.prev_depth != self.curr_depth {
82            if self.prev_depth < self.curr_depth {
83                if let Some(parent_path) = node.path().parent() {
84                    let mut shallow_walker = build_shallow(parent_path, &self.options)?
85                        .into_iter()
86                        .filter_map(|e| e.ok())
87                        .skip_while(|n| n.path() != node.path())
88                        .skip(1);
89
90                    if let Some(_) = shallow_walker.next() {
91                        self.parent_depths.push(self.prev_depth);
92                    } else {
93                        is_last = true;
94                    }
95                }
96            } else {
97                should_pop = true;
98                is_last = true;
99            }
100        }
101
102        display::print(
103            file_name,
104            file_type,
105            self.prev_depth,
106            is_last,
107            &self.parent_depths,
108            self.curr_line_count,
109            self.colours,
110        );
111        if should_pop {
112            self.parent_depths.pop();
113        }
114        Ok(())
115    }
116}