First basic implementation of aggregation; symlinks are not handled yet

3 files changed, 103 insertions, 19 deletions

diff --git a/src/lib.rs b/src/lib.rs
index fd4cb89..5d22577 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,8 +1,69 @@
 extern crate failure;
 extern crate jwalk;
 
-use failure::Error;
+pub struct WalkOptions {
+    pub threads: usize,
+}
+
+impl WalkOptions {
+    pub fn iter_from_path(&self, path: &Path) -> WalkDir {
+        WalkDir::new(path)
+            .preload_metadata(true)
+            .skip_hidden(false)
+            .num_threads(self.threads)
+    }
+}
+
+#[derive(Default)]
+pub struct WalkResult {
+    pub num_errors: usize,
+}
 
-pub fn fun() -> Result<(), Error> {
-    unimplemented!();
+impl fmt::Display for WalkResult {
+    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
+        write!(f, "Encountered {} IO errors", self.num_errors)
+    }
 }
+
+mod aggregate {
+    use crate::{WalkOptions, WalkResult};
+    use failure::Error;
+    use std::io;
+    use std::path::Path;
+
+    pub fn aggregate(
+        mut out: impl io::Write,
+        options: WalkOptions,
+        paths: impl IntoIterator<Item = impl AsRef<Path>>,
+    ) -> Result<WalkResult, Error> {
+        let mut res = WalkResult::default();
+        for path in paths.into_iter() {
+            let mut num_bytes = 0u64;
+            for entry in options.iter_from_path(path.as_ref()) {
+                match entry {
+                    Ok(entry) => {
+                        num_bytes += match entry.metadata {
+                            Some(Ok(m)) => m.len(),
+                            Some(Err(_)) => {
+                                res.num_errors += 1;
+                                0
+                            }
+                            None => unreachable!(
+                                "we ask for metadata, so we at least have Some(Err(..)))"
+                            ),
+                        };
+                    }
+                    Err(_) => res.num_errors += 1,
+                }
+            }
+
+            writeln!(out, "{}\t{}", num_bytes, path.as_ref().display())?;
+        }
+        Ok(res)
+    }
+}
+
+pub use aggregate::aggregate;
+use jwalk::WalkDir;
+use std::fmt;
+use std::path::Path;
diff --git a/src/main.rs b/src/main.rs
index bea5816..73a9f28 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -5,33 +5,55 @@ extern crate structopt;
 
 use failure::Error;
 use failure_tools::ok_or_exit;
+use std::io;
+use std::path::PathBuf;
 use structopt::StructOpt;
 
 mod options {
     use std::path::PathBuf;
 
     #[derive(Debug, StructOpt)]
-    #[structopt(name = "example", about = "An example of StructOpt usage.")]
+    #[structopt(name = "dua", about = "A tool to learn about disk usage, fast!")]
     pub struct Args {
-        /// Activate debug mode
-        #[structopt(short = "d", long = "debug")]
-        debug: bool,
-        /// Set speed
-        #[structopt(short = "s", long = "speed", default_value = "42")]
-        speed: f64,
-        /// Input file
-        #[structopt(parse(from_os_str))]
-        input: PathBuf,
-        /// Output file, stdout if not present
-        #[structopt(parse(from_os_str))]
-        output: Option<PathBuf>,
+        #[structopt(subcommand)]
+        pub command: Option<Command>,
+
+        /// The amount of threads to use. Defaults to the amount of logical processors.
+        /// Set to 1 to use only a single thread.
+        #[structopt(short = "t", long = "threads")]
+        pub threads: Option<usize>,
+    }
+
+    #[derive(Debug, StructOpt)]
+    pub enum Command {
+        #[structopt(name = "aggregate", alias = "a")]
+        Aggregate {
+            /// One or more input files. If unset, we will assume the current directory
+            #[structopt(parse(from_os_str))]
+            input: Vec<PathBuf>,
+        },
     }
 }
 
 fn run() -> Result<(), Error> {
-    let opt = options::Args::from_args();
-    println!("{:?}", opt);
-    dua::fun()
+    use io::Write;
+    use options::Command::*;
+
+    let opt: options::Args = options::Args::from_args();
+    let stdout = io::stdout();
+    let stdout_locked = stdout.lock();
+    let walk_options = dua::WalkOptions {
+        threads: opt.threads.unwrap_or(0),
+    };
+    let res = match opt.command {
+        Some(Aggregate { input: _ }) => unimplemented!(),
+        None => dua::aggregate(stdout_locked, walk_options, vec![PathBuf::from(".")]),
+    }?;
+
+    if res.num_errors > 0 {
+        writeln!(io::stderr(), "{}", res).ok();
+    }
+    Ok(())
 }
 
 fn main() {
diff --git a/tests/snapshots/success-no-arguments b/tests/snapshots/success-no-arguments
new file mode 100644
index 0000000..f8b55b8
--- /dev/null
+++ b/tests/snapshots/success-no-arguments
@@ -0,0 +1 @@
+1258947	.