add command-line command-bytes flags

rggo/firstProgram/wc/main.go

@@ -2,22 +2,37 @@ package main
 
 import (
 	"bufio"
+	"flag"
 	"fmt"
 	"io"
 	"os"
 )
 
 func main() {
-	// Calling the count function to count the number of words
+	// Define a boolean flag -l to count lines instead of words
+	lines := flag.Bool("l", false, "count lines")
+	// Define a boolean flag -b to count bytes instead of words
+	bytes := flag.Bool("b", false, "count bytes")
+
+	// Parse the flags provided by the user
+	flag.Parse()
+	// Calling the count function to count the number of words (or lines)
 	// received from the standard input and printing it out
-	fmt.Println(count(os.Stdin))
+	fmt.Println(count(os.Stdin, *lines, *bytes))
 }
 
-func count(r io.Reader) int {
+func count(r io.Reader, countLines bool, countBytes bool) int {
 	scanner := bufio.NewScanner(r)
-	// Split sets the split function for the Scanner.
-	// The default split function is ScanLines.
-	scanner.Split(bufio.ScanWords)
+	if !countLines {
+		if countBytes {
+			scanner.Split(bufio.ScanBytes)
+		} else {
+			// Split sets the split function for the Scanner.
+			// The default split function is ScanLines.
+			scanner.Split(bufio.ScanWords)
+		}
+	}
+
 	wc := 0
 	// For every word scanned, increment the counter
 	for scanner.Scan() {

rggo/firstProgram/wc/main_test.go

@@ -9,7 +9,27 @@ func TestCountWords(t *testing.T) {
 	b := bytes.NewBufferString("hello world world2 world3 world4\n")
 
 	exp := 5
-	res := count(b)
+	res := count(b, false, false)
+	if res != exp {
+		t.Errorf("Expected %d, got %d instead.\n", exp, res)
+	}
+}
+
+func TestCountLines(t *testing.T) {
+	b := bytes.NewBufferString("hello world world2\nworld3\nworld4\n")
+
+	exp := 3
+	res := count(b, true, false)
+	if res != exp {
+		t.Errorf("Expected %d, got %d instead.\n", exp, res)
+	}
+}
+
+func TestCountBytes(t *testing.T) {
+	b := bytes.NewBufferString("hello\n")
+
+	exp := 6
+	res := count(b, false, true)
 	if res != exp {
 		t.Errorf("Expected %d, got %d instead.\n", exp, res)
 	}