Implement ToSizeInBytes(E)

Allow to convert string with size units (such as "10kb") to bytes.
This converter is privately used in viper. I think it is great to
allow using it from cast. Compared to viper implementation this one
allows to have single-letter suffixes.

Signed-off-by: Evgenii Stratonikov <stratonikov@runbox.com>
This commit is contained in:
Evgenii Stratonikov 2021-10-07 19:35:43 +03:00
parent 88075729b0
commit 3e70cb3062
3 changed files with 102 additions and 0 deletions

View File

@ -25,6 +25,12 @@ func ToTimeInDefaultLocation(i interface{}, location *time.Location) time.Time {
return v return v
} }
// ToSizeInBytes casts an interface to the uint size in bytes.
func ToSizeInBytes(i interface{}) uint {
v, _ := ToSizeInBytesE(i)
return v
}
// ToDuration casts an interface to a time.Duration type. // ToDuration casts an interface to a time.Duration type.
func ToDuration(i interface{}) time.Duration { func ToDuration(i interface{}) time.Duration {
v, _ := ToDurationE(i) v, _ := ToDurationE(i)

View File

@ -9,7 +9,9 @@ import (
"errors" "errors"
"fmt" "fmt"
"html/template" "html/template"
"math"
"path" "path"
"strconv"
"testing" "testing"
"time" "time"
@ -1298,6 +1300,50 @@ func TestToDurationE(t *testing.T) {
} }
} }
func TestToSizeInBytes(t *testing.T) {
const (
kb = 1024
mb = 1024 * kb
gb = 1024 * mb
)
tests := []struct {
input interface{}
expect uint
iserr bool
}{
{123, 123, false},
{ " 42 ", 42, false},
{"1", 1, false},
{"2b", 2, false},
{"3B", 3, false},
{"4 b", 4, false},
{"5k", 5 * kb, false},
{"6 KB", 6 * kb, false},
{"7m", 7 * mb, false},
{"8MB", 8 * mb, false},
{"9 g", 9 * gb, false},
{"10 GB", 10 * gb, false},
// errors
{"test", 0, true},
{strconv.FormatUint(uint64(math.MaxUint64 / kb), 10) + "9999k", 0, true}, // overflow with suffix
{strconv.FormatUint(uint64(math.MaxUint), 10) + "0", 0, true}, // overflow without suffix
}
for i, test := range tests {
errmsg := fmt.Sprintf("i = %d", i) // assert helper message
v, err := ToSizeInBytesE(test.input)
if test.iserr {
assert.Error(t, err, errmsg)
continue
}
assert.NoError(t, err, errmsg)
assert.Equal(t, test.expect, v)
}
}
func TestToTimeWithTimezones(t *testing.T) { func TestToTimeWithTimezones(t *testing.T) {
est, err := time.LoadLocation("EST") est, err := time.LoadLocation("EST")

View File

@ -14,6 +14,7 @@ import (
"strconv" "strconv"
"strings" "strings"
"time" "time"
"unicode"
) )
var errNegativeNotAllowed = errors.New("unable to cast negative value") var errNegativeNotAllowed = errors.New("unable to cast negative value")
@ -51,6 +52,55 @@ func ToTimeInDefaultLocationE(i interface{}, location *time.Location) (tim time.
} }
} }
func safeMul(a, b uint) (uint, error) {
c := a * b
if a > 1 && b > 1 && c/b != a {
return 0, errors.New("overflow")
}
return c, nil
}
// ToSizeInBytesE casts an empty interface to a size in bytes
// interpreting the usual (k, m, g, kb, mB etc.) suffixes.
func ToSizeInBytesE(i interface{}) (uint, error) {
sizeStr := strings.TrimSpace(ToString(i))
lastChar := len(sizeStr) - 1
multiplier := uint(1)
if lastChar > 0 {
if sizeStr[lastChar] == 'b' || sizeStr[lastChar] == 'B' {
sizeStr = sizeStr[:lastChar]
lastChar--
}
if lastChar > 0 {
switch unicode.ToLower(rune(sizeStr[lastChar])) {
case 'k':
multiplier = 1 << 10
sizeStr = strings.TrimSpace(sizeStr[:lastChar])
case 'm':
multiplier = 1 << 20
sizeStr = strings.TrimSpace(sizeStr[:lastChar])
case 'g':
multiplier = 1 << 30
sizeStr = strings.TrimSpace(sizeStr[:lastChar])
default:
multiplier = 1
sizeStr = strings.TrimSpace(sizeStr[:lastChar+1])
}
}
}
size, err := ToUintE(sizeStr)
if err != nil {
return 0, err
}
return safeMul(size, multiplier)
}
// ToDurationE casts an interface to a time.Duration type. // ToDurationE casts an interface to a time.Duration type.
func ToDurationE(i interface{}) (d time.Duration, err error) { func ToDurationE(i interface{}) (d time.Duration, err error) {
i = indirect(i) i = indirect(i)