Support for various byte formats

5 files changed, 93 insertions, 29 deletions

diff --git a/src/lib.rs b/src/lib.rs
index ee3b35a..65983cc 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,33 +1,51 @@
 extern crate failure;
 extern crate jwalk;
 
-pub struct WalkOptions {
-    pub threads: usize,
-}
+mod common {
+    use jwalk::WalkDir;
+    use std::{fmt, path::Path};
 
-impl WalkOptions {
-    pub fn format_bytes(&self, b: u64) -> String {
-        use byte_unit::Byte;
-        Byte::from_bytes(b as u128)
-            .get_appropriate_unit(false)
-            .format(2)
+    pub enum ByteFormat {
+        Metric,
+        Binary,
+        Bytes,
     }
-    pub fn iter_from_path(&self, path: &Path) -> WalkDir {
-        WalkDir::new(path)
-            .preload_metadata(true)
-            .skip_hidden(false)
-            .num_threads(self.threads)
+
+    pub struct WalkOptions {
+        pub threads: usize,
+        pub format: ByteFormat,
     }
-}
 
-#[derive(Default)]
-pub struct WalkResult {
-    pub num_errors: usize,
-}
+    impl WalkOptions {
+        pub fn format_bytes(&self, b: u64) -> String {
+            use byte_unit::Byte;
+            use ByteFormat::*;
+            let binary = match self.format {
+                Bytes => return format!("{} b", b),
+                Binary => true,
+                Metric => false,
+            };
+            Byte::from_bytes(b as u128)
+                .get_appropriate_unit(binary)
+                .format(2)
+        }
+        pub fn iter_from_path(&self, path: &Path) -> WalkDir {
+            WalkDir::new(path)
+                .preload_metadata(true)
+                .skip_hidden(false)
+                .num_threads(self.threads)
+        }
+    }
 
-impl fmt::Display for WalkResult {
-    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
-        write!(f, "Encountered {} IO errors", self.num_errors)
+    #[derive(Default)]
+    pub struct WalkResult {
+        pub num_errors: usize,
+    }
+
+    impl fmt::Display for WalkResult {
+        fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
+            write!(f, "Encountered {} IO errors", self.num_errors)
+        }
     }
 }
 
@@ -55,7 +73,7 @@ mod aggregate {
                                 0
                             }
                             None => unreachable!(
-                                "we ask for metadata, so we at least have Some(Err(..)))"
+                                "we ask for metadata, so we at least have Some(Err(..))). Issue in jwalk?"
                             ),
                         };
                     }
@@ -75,5 +93,4 @@ mod aggregate {
 }
 
 pub use aggregate::aggregate;
-use jwalk::WalkDir;
-use std::{fmt, path::Path};
+pub use common::*;
diff --git a/src/main.rs b/src/main.rs
index 11f0670..89951dc 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,15 +1,37 @@
 extern crate failure;
 extern crate failure_tools;
-#[macro_use]
 extern crate structopt;
 
+use structopt::StructOpt;
+
+use dua::ByteFormat;
 use failure::Error;
 use failure_tools::ok_or_exit;
 use std::{io, path::PathBuf};
-use structopt::StructOpt;
 
 mod options {
+    use dua::ByteFormat as LibraryByteFormat;
     use std::path::PathBuf;
+    use structopt::{clap::arg_enum, StructOpt};
+
+    arg_enum! {
+        #[derive(PartialEq, Debug)]
+        pub enum ByteFormat {
+            HumanMetric,
+            HumanBinary,
+            Bytes
+        }
+    }
+
+    impl From<ByteFormat> for LibraryByteFormat {
+        fn from(input: ByteFormat) -> Self {
+            match input {
+                ByteFormat::HumanMetric => LibraryByteFormat::Metric,
+                ByteFormat::HumanBinary => LibraryByteFormat::Binary,
+                ByteFormat::Bytes => LibraryByteFormat::Bytes,
+            }
+        }
+    }
 
     #[derive(Debug, StructOpt)]
     #[structopt(name = "dua", about = "A tool to learn about disk usage, fast!")]
@@ -21,6 +43,13 @@ mod options {
         /// Set to 1 to use only a single thread.
         #[structopt(short = "t", long = "threads")]
         pub threads: Option<usize>,
+
+        /// The format with which to print byte counts.
+        /// HumanMetric - uses 1000 as base (default)
+        /// HumanBinary - uses 1024 as base
+        /// Bytes - plain bytes without any formatting
+        #[structopt(short = "f", long = "format")]
+        pub format: Option<ByteFormat>,
     }
 
     #[derive(Debug, StructOpt)]
@@ -43,6 +72,7 @@ fn run() -> Result<(), Error> {
     let stdout_locked = stdout.lock();
     let walk_options = dua::WalkOptions {
         threads: opt.threads.unwrap_or(0),
+        format: opt.format.map(Into::into).unwrap_or(ByteFormat::Metric),
     };
     let res = match opt.command {
         Some(Aggregate { input: _ }) => unimplemented!(),
diff --git a/tests/snapshots/success-bytes-binary b/tests/snapshots/success-bytes-binary
new file mode 100644
index 0000000..59210ca
--- /dev/null
+++ b/tests/snapshots/success-bytes-binary
@@ -0,0 +1 @@
+1.20 MiB	.
\ No newline at end of file
diff --git a/tests/snapshots/success-bytes-only b/tests/snapshots/success-bytes-only
new file mode 100644
index 0000000..e0c2326
--- /dev/null
+++ b/tests/snapshots/success-bytes-only
@@ -0,0 +1 @@
+1258947 b	.
\ No newline at end of file
diff --git a/tests/stateless-journey.sh b/tests/stateless-journey.sh
index b170af3..d692b14 100755
--- a/tests/stateless-journey.sh
+++ b/tests/stateless-journey.sh
@@ -15,11 +15,26 @@ SUCCESSFULLY=0
 (with "a sample directory"
   (sandbox
     cp -R "$fixtures/sample-01" .
-    (when "running the program without arguments"
-      it "produces a human-readable aggregate of the current directory" && {
+    (with "no arguments"
+      it "produces a human-readable (metric) aggregate of the current directory, without total" && {
         WITH_SNAPSHOT="$snapshot/success-no-arguments" \
         expect_run ${SUCCESSFULLY} "$exe"
       }
     )
+
+    (with "the byte format set"
+      (with "human-binary"
+        it "produces a human-readable aggregate of the current directory, without total" && {
+          WITH_SNAPSHOT="$snapshot/success-bytes-binary" \
+          expect_run ${SUCCESSFULLY} "$exe" --format humanbinary
+        }
+      )
+      (with "bytes"
+        it "produces a human-readable aggregate of the current directory, without total" && {
+          WITH_SNAPSHOT="$snapshot/success-bytes-only" \
+          expect_run ${SUCCESSFULLY} "$exe" --format bytes
+        }
+      )
+    )
   )
 )