pkg: add new datasizes package that adds constants and parsing

This commit extracts the existing `common.DataSizeToUint64` and the `common.{Ki,Gi,Mi,...}` and similar data size constants into a new package `datasizes` that will be shared with other packages like `blueprints`.
osbuild · Oct 17, 2024 · bdef54d · bdef54d
1 parent 7b89d25
commit bdef54d
Show file tree

Hide file tree

Showing 6 changed files with 123 additions and 107 deletions.
diff --git a/internal/common/constants.go b/internal/common/constants.go
@@ -2,25 +2,6 @@ package common
 
 import "fmt"
 
-const (
-	KiloByte = 1000                      // kB
-	KibiByte = 1024                      // KiB
-	MegaByte = 1000 * 1000               // MB
-	MebiByte = 1024 * 1024               // MiB
-	GigaByte = 1000 * 1000 * 1000        // GB
-	GibiByte = 1024 * 1024 * 1024        // GiB
-	TeraByte = 1000 * 1000 * 1000 * 1000 // TB
-	TebiByte = 1024 * 1024 * 1024 * 1024 // TiB
-
-	// shorthands
-	KiB = KibiByte
-	MB  = MegaByte
-	MiB = MebiByte
-	GB  = GigaByte
-	GiB = GibiByte
-	TiB = TebiByte
-)
-
 // These constants are set during buildtime using additional
 // compiler flags. Not all of them are necessarily defined
 // because RPMs can be build from a tarball and spec file without

diff --git a/internal/common/helpers.go b/internal/common/helpers.go
@@ -5,9 +5,7 @@ import (
 	"fmt"
 	"io"
 	"os/exec"
-	"regexp"
 	"sort"
-	"strconv"
 	"strings"
 )
 
@@ -27,54 +25,6 @@ func IsStringInSortedSlice(slice []string, s string) bool {
 	return false
 }
 
-// DataSizeToUint64 converts a size specified as a string in KB/KiB/MB/etc. to
-// a number of bytes represented by uint64.
-func DataSizeToUint64(size string) (uint64, error) {
-	// Pre-process the input
-	size = strings.TrimSpace(size)
-
-	// Get the number from the string
-	plain_number := regexp.MustCompile(`[[:digit:]]+`)
-	number_as_str := plain_number.FindString(size)
-	if number_as_str == "" {
-		return 0, fmt.Errorf("the size string doesn't contain any number: %s", size)
-	}
-
-	// Parse the number into integer
-	return_size, err := strconv.ParseUint(number_as_str, 10, 64)
-	if err != nil {
-		return 0, fmt.Errorf("failed to parse size as integer: %s", number_as_str)
-	}
-
-	// List of all supported units (from kB to TB and KiB to TiB)
-	supported_units := []struct {
-		re       *regexp.Regexp
-		multiple uint64
-	}{
-		{regexp.MustCompile(`^\s*[[:digit:]]+\s*kB$`), KiloByte},
-		{regexp.MustCompile(`^\s*[[:digit:]]+\s*KiB$`), KibiByte},
-		{regexp.MustCompile(`^\s*[[:digit:]]+\s*MB$`), MegaByte},
-		{regexp.MustCompile(`^\s*[[:digit:]]+\s*MiB$`), MebiByte},
-		{regexp.MustCompile(`^\s*[[:digit:]]+\s*GB$`), GigaByte},
-		{regexp.MustCompile(`^\s*[[:digit:]]+\s*GiB$`), GibiByte},
-		{regexp.MustCompile(`^\s*[[:digit:]]+\s*TB$`), TeraByte},
-		{regexp.MustCompile(`^\s*[[:digit:]]+\s*TiB$`), TebiByte},
-		{regexp.MustCompile(`^\s*[[:digit:]]+$`), 1},
-	}
-
-	for _, unit := range supported_units {
-		if unit.re.MatchString(size) {
-			return_size *= unit.multiple
-			return return_size, nil
-		}
-	}
-
-	// In case the strign didn't match any of the above regexes, return nil
-	// even if a number was found. This is to prevent users from submitting
-	// unknown units.
-	return 0, fmt.Errorf("unknown data size units in string: %s", size)
-}
-
 // NopSeekCloser returns an io.ReadSeekCloser with a no-op Close method
 // wrapping the provided io.ReadSeeker r.
 func NopSeekCloser(r io.ReadSeeker) io.ReadSeekCloser {

diff --git a/internal/common/helpers_test.go b/internal/common/helpers_test.go
@@ -5,7 +5,6 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
 )
 
 func TestPanicOnError(t *testing.T) {
@@ -20,43 +19,6 @@ func TestIsStringInSortedSlice(t *testing.T) {
 	assert.False(t, IsStringInSortedSlice([]string{}, "homer"))
 }
 
-func TestDataSizeToUint64(t *testing.T) {
-	cases := []struct {
-		input   string
-		success bool
-		output  uint64
-	}{
-		{"123", true, 123},
-		{"123 kB", true, 123000},
-		{"123 KiB", true, 123 * 1024},
-		{"123 MB", true, 123 * 1000 * 1000},
-		{"123 MiB", true, 123 * 1024 * 1024},
-		{"123 GB", true, 123 * 1000 * 1000 * 1000},
-		{"123 GiB", true, 123 * 1024 * 1024 * 1024},
-		{"123 TB", true, 123 * 1000 * 1000 * 1000 * 1000},
-		{"123 TiB", true, 123 * 1024 * 1024 * 1024 * 1024},
-		{"123kB", true, 123000},
-		{"123KiB", true, 123 * 1024},
-		{" 123  ", true, 123},
-		{"  123kB  ", true, 123000},
-		{"  123KiB  ", true, 123 * 1024},
-		{"123 KB", false, 0},
-		{"123 mb", false, 0},
-		{"123 PB", false, 0},
-		{"123 PiB", false, 0},
-	}
-
-	for _, c := range cases {
-		result, err := DataSizeToUint64(c.input)
-		if c.success {
-			require.Nil(t, err)
-			assert.EqualValues(t, c.output, result)
-		} else {
-			assert.NotNil(t, err)
-		}
-	}
-}
-
 func TestSystemdMountUnit(t *testing.T) {
 	for _, tc := range []struct {
 		mountpoint   string

diff --git a/pkg/datasizes/constants.go b/pkg/datasizes/constants.go
@@ -0,0 +1,20 @@
+package datasizes
+
+const (
+	KiloByte = 1000                      // kB
+	KibiByte = 1024                      // KiB
+	MegaByte = 1000 * 1000               // MB
+	MebiByte = 1024 * 1024               // MiB
+	GigaByte = 1000 * 1000 * 1000        // GB
+	GibiByte = 1024 * 1024 * 1024        // GiB
+	TeraByte = 1000 * 1000 * 1000 * 1000 // TB
+	TebiByte = 1024 * 1024 * 1024 * 1024 // TiB
+
+	// shorthands
+	KiB = KibiByte
+	MB  = MegaByte
+	MiB = MebiByte
+	GB  = GigaByte
+	GiB = GibiByte
+	TiB = TebiByte
+)
diff --git a/pkg/datasizes/parse.go b/pkg/datasizes/parse.go
@@ -0,0 +1,56 @@
+package datasizes
+
+import (
+	"fmt"
+	"regexp"
+	"strconv"
+	"strings"
+)
+
+// Parse converts a size specified as a string in KB/KiB/MB/etc. to
+// a number of bytes represented by uint64.
+func Parse(size string) (uint64, error) {
+	// Pre-process the input
+	size = strings.TrimSpace(size)
+
+	// Get the number from the string
+	plain_number := regexp.MustCompile(`[[:digit:]]+`)
+	number_as_str := plain_number.FindString(size)
+	if number_as_str == "" {
+		return 0, fmt.Errorf("the size string doesn't contain any number: %s", size)
+	}
+
+	// Parse the number into integer
+	return_size, err := strconv.ParseUint(number_as_str, 10, 64)
+	if err != nil {
+		return 0, fmt.Errorf("failed to parse size as integer: %s", number_as_str)
+	}
+
+	// List of all supported units (from kB to TB and KiB to TiB)
+	supported_units := []struct {
+		re       *regexp.Regexp
+		multiple uint64
+	}{
+		{regexp.MustCompile(`^\s*[[:digit:]]+\s*kB$`), KiloByte},
+		{regexp.MustCompile(`^\s*[[:digit:]]+\s*KiB$`), KibiByte},
+		{regexp.MustCompile(`^\s*[[:digit:]]+\s*MB$`), MegaByte},
+		{regexp.MustCompile(`^\s*[[:digit:]]+\s*MiB$`), MebiByte},
+		{regexp.MustCompile(`^\s*[[:digit:]]+\s*GB$`), GigaByte},
+		{regexp.MustCompile(`^\s*[[:digit:]]+\s*GiB$`), GibiByte},
+		{regexp.MustCompile(`^\s*[[:digit:]]+\s*TB$`), TeraByte},
+		{regexp.MustCompile(`^\s*[[:digit:]]+\s*TiB$`), TebiByte},
+		{regexp.MustCompile(`^\s*[[:digit:]]+$`), 1},
+	}
+
+	for _, unit := range supported_units {
+		if unit.re.MatchString(size) {
+			return_size *= unit.multiple
+			return return_size, nil
+		}
+	}
+
+	// In case the strign didn't match any of the above regexes, return nil
+	// even if a number was found. This is to prevent users from submitting
+	// unknown units.
+	return 0, fmt.Errorf("unknown data size units in string: %s", size)
+}
diff --git a/pkg/datasizes/parse_test.go b/pkg/datasizes/parse_test.go
@@ -0,0 +1,47 @@
+package datasizes_test
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/osbuild/images/pkg/datasizes"
+)
+
+func TestDataSizeToUint64(t *testing.T) {
+	cases := []struct {
+		input   string
+		success bool
+		output  uint64
+	}{
+		{"123", true, 123},
+		{"123 kB", true, 123000},
+		{"123 KiB", true, 123 * 1024},
+		{"123 MB", true, 123 * 1000 * 1000},
+		{"123 MiB", true, 123 * 1024 * 1024},
+		{"123 GB", true, 123 * 1000 * 1000 * 1000},
+		{"123 GiB", true, 123 * 1024 * 1024 * 1024},
+		{"123 TB", true, 123 * 1000 * 1000 * 1000 * 1000},
+		{"123 TiB", true, 123 * 1024 * 1024 * 1024 * 1024},
+		{"123kB", true, 123000},
+		{"123KiB", true, 123 * 1024},
+		{" 123  ", true, 123},
+		{"  123kB  ", true, 123000},
+		{"  123KiB  ", true, 123 * 1024},
+		{"123 KB", false, 0},
+		{"123 mb", false, 0},
+		{"123 PB", false, 0},
+		{"123 PiB", false, 0},
+	}
+
+	for _, c := range cases {
+		result, err := datasizes.Parse(c.input)
+		if c.success {
+			require.Nil(t, err)
+			assert.EqualValues(t, c.output, result)
+		} else {
+			assert.NotNil(t, err)
+		}
+	}
+}