diff --git a/cmd/randomtext/main.go b/cmd/randomtext/main.go index 055e294..38c8f42 100644 --- a/cmd/randomtext/main.go +++ b/cmd/randomtext/main.go @@ -1,6 +1,7 @@ package main import ( + "errors" "flag" "fmt" "github.com/kishaningithub/randomtext" @@ -14,38 +15,50 @@ func main() { sizePtr := flag.String("size", "1MB", "Size of generated random text in KB, MB, GB, TB") typePtr := flag.String("type", "chars", "Type of text to be generated - chars, words, zeros") flag.Parse() - sizeInBytes := parseSize(strings.ToLower(*sizePtr)) - generator := parseType(strings.ToLower(*typePtr)) + sizeInBytes, err := parseSize(strings.ToLower(*sizePtr)) + handleErr(err) + generator, err := parseType(strings.ToLower(*typePtr)) + handleErr(err) randomtext.Generate(sizeInBytes, generator, os.Stdout) } -func parseSize(size string) int { +func handleErr(err error) { + if err != nil { + fmt.Println(err) + os.Exit(1) + } +} + +func parseSize(size string) (int, error) { reg := regexp.MustCompile("\\d+") - givenSize, _ := strconv.Atoi(reg.FindString(size)) - sizeInBytes := 1 + givenSize, err := strconv.Atoi(reg.FindString(size)) + if err != nil { + return 0, errors.New("invalid size") + } if strings.Contains(size, "kb") { - sizeInBytes = givenSize * 1024 - } else if strings.Contains(size, "mb") { - sizeInBytes = givenSize * 1024 * 1024 - } else if strings.Contains(size, "gb") { - sizeInBytes = givenSize * 1024 * 1024 * 1024 - } else if strings.Contains(size, "tb") { - sizeInBytes = givenSize * 1024 * 1024 * 1024 * 1024 + return givenSize * 1024, nil } - return sizeInBytes + if strings.Contains(size, "mb") { + return givenSize * 1024 * 1024, nil + } + if strings.Contains(size, "gb") { + return givenSize * 1024 * 1024 * 1024, nil + } + if strings.Contains(size, "tb") { + return givenSize * 1024 * 1024 * 1024 * 1024, nil + } + return givenSize, nil } -func parseType(typeStr string) func() string { +func parseType(typeStr string) (func() string, error) { switch typeStr { case "chars": - return randomtext.CharGenerator() + return randomtext.CharGenerator(), nil case "words": - return randomtext.WordGenerator() + return randomtext.WordGenerator(), nil case "zeros": - return randomtext.ZeroGenerator() + return randomtext.ZeroGenerator(), nil default: - fmt.Print("Not a valid value") - os.Exit(1) + return nil, errors.New("not a valid value") } - return nil }