pkg/units: including suggestions and enhancements

Docker-DCO-1.1-Signed-off-by: fcarriedo <fcarriedo@gmail.com> (github: fcarriedo)

pkg/units/size.go

@@ -7,9 +7,30 @@ import (
 	"strings"
 )
 
+type unit int64
+
+// See: http://en.wikipedia.org/wiki/Binary_prefix
 const (
-	decimalKUnit = 1000
-	binaryKUnit  = 1024
+	// Decimal
+	KB unit = 1000
+	MB      = 1000 * KB
+	GB      = 1000 * MB
+	TB      = 1000 * GB
+	PB      = 1000 * TB
+
+	// Binary
+	KiB unit = 1024
+	MiB      = 1024 * KiB
+	GiB      = 1024 * MiB
+	TiB      = 1024 * GiB
+	PiB      = 1024 * TiB
+)
+
+type unitMap map[string]unit
+
+var (
+	decimalMap = unitMap{"k": KB, "m": MB, "g": GB, "t": TB, "p": PB}
+	binaryMap  = unitMap{"k": KiB, "m": MiB, "g": GiB, "t": TiB, "p": PiB}
 )
 
 var sizeRegex *regexp.Regexp
@@ -18,7 +39,7 @@ func init() {
 	sizeRegex = regexp.MustCompile("^(\\d+)([kKmMgGtTpP])?[bB]?$")
 }
 
-var bytePrefixes = [...]string{"B", "kB", "MB", "GB", "TB", "PB", "EB", "ZB", "YB"}
+var unitAbbrs = [...]string{"B", "kB", "MB", "GB", "TB", "PB", "EB", "ZB", "YB"}
 
 // HumanSize returns a human-readable approximation of a size
 // using SI standard (eg. "44kB", "17MB")
@@ -29,13 +50,13 @@ func HumanSize(size int64) string {
 		sizef = sizef / 1000.0
 		i++
 	}
-	return fmt.Sprintf("%.4g %s", sizef, bytePrefixes[i])
+	return fmt.Sprintf("%.4g %s", sizef, unitAbbrs[i])
 }
 
 // FromHumanSize returns an integer from a human-readable specification of a
 // size using SI standard (eg. "44kB", "17MB")
 func FromHumanSize(size string) (int64, error) {
-	return parseSize(size, decimalKUnit)
+	return parseSize(size, decimalMap)
 }
 
 // Parses a human-readable string representing an amount of RAM
@@ -43,37 +64,25 @@ func FromHumanSize(size string) (int64, error) {
 // returns the number of bytes, or -1 if the string is unparseable.
 // Units are case-insensitive, and the 'b' suffix is optional.
 func RAMInBytes(size string) (int64, error) {
-	return parseSize(size, binaryKUnit)
+	return parseSize(size, binaryMap)
 }
 
-// Parses the human-readable size string into the amount it represents given
-// the desired kilo unit [decimalKiloUnit=1000|binaryKiloUnit=1024]
-func parseSize(size string, kUnit int64) (int64, error) {
-	matches := sizeRegex.FindStringSubmatch(size)
-
+// Parses the human-readable size string into the amount it represents
+func parseSize(sizeStr string, uMap unitMap) (int64, error) {
+	matches := sizeRegex.FindStringSubmatch(sizeStr)
 	if len(matches) != 3 {
-		return -1, fmt.Errorf("Invalid size: '%s'", size)
+		return -1, fmt.Errorf("Invalid size: '%s'", sizeStr)
 	}
 
-	theSize, err := strconv.ParseInt(matches[1], 10, 0)
+	size, err := strconv.ParseInt(matches[1], 10, 0)
 	if err != nil {
 		return -1, err
 	}
 
 	unitPrefix := strings.ToLower(matches[2])
-
-	switch unitPrefix {
-	case "k":
-		theSize *= kUnit
-	case "m":
-		theSize *= kUnit * kUnit
-	case "g":
-		theSize *= kUnit * kUnit * kUnit
-	case "t":
-		theSize *= kUnit * kUnit * kUnit * kUnit
-	case "p":
-		theSize *= kUnit * kUnit * kUnit * kUnit * kUnit
+	if mul, ok := uMap[unitPrefix]; ok {
+		size *= int64(mul)
 	}
 
-	return theSize, nil
+	return size, nil
 }