cmds/core/du: add -k flag

Signed-off-by: Siarhiej Siemianczuk <[email protected]>
archie2x · Aug 29, 2024 · 3546135 · 3546135
1 parent 4533532
commit 3546135
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 10 deletions.
diff --git a/cmds/core/du/du.go b/cmds/core/du/du.go
@@ -23,12 +23,14 @@ var errNoStatInfo = errors.New("os.FileInfo has no stat_t info")
 type cmd struct {
 	stdout      io.Writer
 	reportFiles bool
+	kbUnit      bool
 }
 
-func command(stdout io.Writer, reportFiles bool) *cmd {
+func command(stdout io.Writer, reportFiles, kbUnit bool) *cmd {
 	return &cmd{
 		stdout:      stdout,
 		reportFiles: reportFiles,
+		kbUnit:      kbUnit,
 	}
 }
 
@@ -42,7 +44,7 @@ func (c *cmd) run(files ...string) error {
 		if err != nil {
 			return err
 		}
-		fmt.Fprintf(c.stdout, "%d\t%s\n", blocks, file)
+		c.print(blocks, file)
 	}
 
 	return nil
@@ -65,7 +67,7 @@ func (c *cmd) du(file string) (int64, error) {
 
 			blocks += dirBlocks
 
-			fmt.Fprintf(c.stdout, "%d\t%s\n", dirBlocks, path)
+			c.print(dirBlocks, path)
 			return fs.SkipDir
 		}
 
@@ -75,7 +77,7 @@ func (c *cmd) du(file string) (int64, error) {
 		}
 
 		if c.reportFiles && !info.IsDir() {
-			fmt.Fprintf(c.stdout, "%d\t%s\n", st.Blocks, path)
+			c.print(st.Blocks, path)
 		}
 
 		blocks += st.Blocks
@@ -85,10 +87,18 @@ func (c *cmd) du(file string) (int64, error) {
 	return blocks, nil
 }
 
+func (c *cmd) print(nblock int64, path string) {
+	if c.kbUnit {
+		nblock /= 2
+	}
+	fmt.Fprintf(c.stdout, "%d\t%s\n", nblock, path)
+}
+
 func main() {
 	var reportFiles = flag.Bool("a", false, "report the size of each file not of type directory")
+	var kbUnit = flag.Bool("k", false, "write the files sizes in units of 1024 bytes, rather than the default 512-byte units")
 	flag.Parse()
-	if err := command(os.Stdout, *reportFiles).run(flag.Args()...); err != nil {
+	if err := command(os.Stdout, *reportFiles, *kbUnit).run(flag.Args()...); err != nil {
 		log.Fatalf("du: %v", err)
 	}
 }
diff --git a/cmds/core/du/du_test.go b/cmds/core/du/du_test.go
@@ -22,7 +22,7 @@ func TestDU(t *testing.T) {
 		}
 		f.Write(make([]byte, 8096))
 
-		blocks, err := command(io.Discard, false).du(f.Name())
+		blocks, err := command(io.Discard, false, false).du(f.Name())
 		if err != nil {
 			t.Fatalf("expected nil got %v", err)
 		}
@@ -38,7 +38,7 @@ func TestDU(t *testing.T) {
 			t.Fatal(err)
 		}
 
-		blocks, err := command(io.Discard, false).du(f.Name())
+		blocks, err := command(io.Discard, false, false).du(f.Name())
 		if err != nil {
 			t.Fatalf("expected nil got %v", err)
 		}
@@ -55,7 +55,7 @@ func TestDU(t *testing.T) {
 		}
 		f.Write(make([]byte, 1))
 
-		blocks, err := command(io.Discard, false).du(f.Name())
+		blocks, err := command(io.Discard, false, false).du(f.Name())
 		if err != nil {
 			t.Fatalf("expected nil got %v", err)
 		}
@@ -75,7 +75,7 @@ func TestRun(t *testing.T) {
 		}
 
 		stdout := &bytes.Buffer{}
-		err = command(stdout, false).run()
+		err = command(stdout, false, false).run()
 		if err != nil {
 			t.Fatalf("expected nil got %v", err)
 		}
@@ -112,7 +112,7 @@ func TestRun(t *testing.T) {
 		}
 
 		stdout := &bytes.Buffer{}
-		err = command(stdout, true).run(dir)
+		err = command(stdout, true, false).run(dir)
 		if err != nil {
 			t.Fatalf("expected nil got %v", err)
 		}
@@ -123,4 +123,22 @@ func TestRun(t *testing.T) {
 			t.Errorf("expected file1, file2 and temp dir, but got %d lines", len(lines))
 		}
 	})
+	t.Run("with -k flag", func(t *testing.T) {
+		dir := t.TempDir()
+		f, err := os.CreateTemp(dir, "")
+		if err != nil {
+			t.Fatal(err)
+		}
+		f.Write(make([]byte, 4096))
+
+		stdout := &bytes.Buffer{}
+		err = command(stdout, false, true).run(f.Name())
+		if err != nil {
+			t.Fatalf("expected nil got %v", err)
+		}
+
+		if stdout.String()[0] != '4' {
+			t.Errorf("expected 4 blocks with -k, got %q", stdout.String())
+		}
+	})
 }