parallel_disk_usage/app/
sub.rs

1use crate::{
2    args::Fraction,
3    data_tree::{DataTree, DataTreeReflection},
4    fs_tree_builder::FsTreeBuilder,
5    get_size::GetSize,
6    json_data::{BinaryVersion, JsonData, SchemaVersion, UnitAndTree},
7    os_string_display::OsStringDisplay,
8    reporter::ParallelReporter,
9    runtime_error::RuntimeError,
10    size,
11    status_board::GLOBAL_STATUS_BOARD,
12    visualizer::{BarAlignment, ColumnWidthDistribution, Direction, Visualizer},
13};
14use serde::Serialize;
15use std::{io::stdout, iter::once, num::NonZeroUsize, path::PathBuf};
16
17/// The sub program of the main application.
18pub struct Sub<Size, SizeGetter, Report>
19where
20    Report: ParallelReporter<Size> + Sync,
21    Size: size::Size + Into<u64> + Serialize + Send + Sync,
22    SizeGetter: GetSize<Size = Size> + Copy + Sync,
23    DataTreeReflection<String, Size>: Into<UnitAndTree>,
24{
25    /// List of files and/or directories.
26    pub files: Vec<PathBuf>,
27    /// Print JSON data instead of an ASCII chart.
28    pub json_output: bool,
29    /// Format to be used to [`display`](size::Size::display) the sizes returned by [`size_getter`](Self::size_getter).
30    pub bytes_format: Size::DisplayFormat,
31    /// The direction of the visualization.
32    pub direction: Direction,
33    /// The alignment of the bars.
34    pub bar_alignment: BarAlignment,
35    /// Distribution and number of characters/blocks can be placed in a line.
36    pub column_width_distribution: ColumnWidthDistribution,
37    /// Maximum number of levels that should be visualized.
38    pub max_depth: NonZeroUsize,
39    /// [Get the size](GetSize) of files/directories.
40    pub size_getter: SizeGetter,
41    /// Reports measurement progress.
42    pub reporter: Report,
43    /// Minimal size proportion required to appear.
44    pub min_ratio: Fraction,
45    /// Preserve order of entries.
46    pub no_sort: bool,
47}
48
49impl<Size, SizeGetter, Report> Sub<Size, SizeGetter, Report>
50where
51    Size: size::Size + Into<u64> + Serialize + Send + Sync,
52    Report: ParallelReporter<Size> + Sync,
53    SizeGetter: GetSize<Size = Size> + Copy + Sync,
54    DataTreeReflection<String, Size>: Into<UnitAndTree>,
55{
56    /// Run the sub program.
57    pub fn run(self) -> Result<(), RuntimeError> {
58        let Sub {
59            files,
60            json_output,
61            bytes_format,
62            direction,
63            bar_alignment,
64            column_width_distribution,
65            max_depth,
66            size_getter,
67            reporter,
68            min_ratio,
69            no_sort,
70        } = self;
71
72        let mut iter = files
73            .into_iter()
74            .map(|root| -> DataTree<OsStringDisplay, Size> {
75                FsTreeBuilder {
76                    reporter: &reporter,
77                    root,
78                    size_getter,
79                }
80                .into()
81            });
82
83        let data_tree = if let Some(data_tree) = iter.next() {
84            data_tree
85        } else {
86            return Sub {
87                files: vec![".".into()],
88                reporter,
89                ..self
90            }
91            .run();
92        };
93
94        // ExactSizeIterator::is_empty is unstable
95        let data_tree = if iter.len() == 0 {
96            data_tree
97        } else {
98            let children: Vec<_> = once(data_tree).chain(iter).collect();
99            DataTree::dir(
100                OsStringDisplay::os_string_from("(total)"),
101                Size::default(),
102                children,
103            )
104        };
105
106        if reporter.destroy().is_err() {
107            eprintln!("[warning] Failed to destroy the thread that reports progress");
108        }
109
110        let min_ratio: f32 = min_ratio.into();
111        let data_tree = {
112            let mut data_tree = data_tree;
113            if min_ratio > 0.0 {
114                data_tree.par_cull_insignificant_data(min_ratio);
115            }
116            if !no_sort {
117                data_tree.par_sort_by(|left, right| left.size().cmp(&right.size()).reverse());
118            }
119            data_tree
120        };
121
122        GLOBAL_STATUS_BOARD.clear_line(0);
123
124        if json_output {
125            let unit_and_tree: UnitAndTree = data_tree
126                .into_reflection() // I really want to use std::mem::transmute here but can't.
127                .par_convert_names_to_utf8() // TODO: allow non-UTF8 somehow.
128                .expect("convert all names from raw string to UTF-8")
129                .into();
130            let json_data = JsonData {
131                schema_version: SchemaVersion,
132                binary_version: Some(BinaryVersion::current()),
133                unit_and_tree,
134            };
135            return serde_json::to_writer(stdout(), &json_data)
136                .map_err(RuntimeError::SerializationFailure);
137        }
138
139        let visualizer = Visualizer {
140            data_tree: &data_tree,
141            bytes_format,
142            direction,
143            bar_alignment,
144            column_width_distribution,
145            max_depth,
146        };
147
148        print!("{visualizer}"); // visualizer already ends with "\n", println! isn't needed here.
149        Ok(())
150    }
151}