forked from mirror/viper
Add GetSizeInBytes.
Useful to parse strings like 1GB or 12 mb into an unsigned integer number of bytes.
This commit is contained in:
parent
5b0b926e3d
commit
e133904c4f
45
util.go
45
util.go
|
@ -19,8 +19,10 @@ import (
|
||||||
"path/filepath"
|
"path/filepath"
|
||||||
"runtime"
|
"runtime"
|
||||||
"strings"
|
"strings"
|
||||||
|
"unicode"
|
||||||
|
|
||||||
"github.com/BurntSushi/toml"
|
"github.com/BurntSushi/toml"
|
||||||
|
"github.com/spf13/cast"
|
||||||
jww "github.com/spf13/jwalterweatherman"
|
jww "github.com/spf13/jwalterweatherman"
|
||||||
"gopkg.in/yaml.v2"
|
"gopkg.in/yaml.v2"
|
||||||
)
|
)
|
||||||
|
@ -139,3 +141,46 @@ func marshallConfigReader(in io.Reader, c map[string]interface{}, configType str
|
||||||
|
|
||||||
insensitiviseMap(c)
|
insensitiviseMap(c)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func safeMul(a, b uint) uint {
|
||||||
|
c := a * b
|
||||||
|
if a > 1 && b > 1 && c/b != a {
|
||||||
|
return 0
|
||||||
|
}
|
||||||
|
return c
|
||||||
|
}
|
||||||
|
|
||||||
|
// parseSizeInBytes converts strings like 1GB or 12 mb into an unsigned integer number of bytes
|
||||||
|
func parseSizeInBytes(sizeStr string) uint {
|
||||||
|
sizeStr = strings.TrimSpace(sizeStr)
|
||||||
|
lastChar := len(sizeStr) - 1
|
||||||
|
multiplier := uint(1)
|
||||||
|
|
||||||
|
if lastChar > 0 {
|
||||||
|
if sizeStr[lastChar] == 'b' || sizeStr[lastChar] == 'B' {
|
||||||
|
if lastChar > 1 {
|
||||||
|
switch unicode.ToLower(rune(sizeStr[lastChar-1])) {
|
||||||
|
case 'k':
|
||||||
|
multiplier = 1 << 10
|
||||||
|
sizeStr = strings.TrimSpace(sizeStr[:lastChar-1])
|
||||||
|
case 'm':
|
||||||
|
multiplier = 1 << 20
|
||||||
|
sizeStr = strings.TrimSpace(sizeStr[:lastChar-1])
|
||||||
|
case 'g':
|
||||||
|
multiplier = 1 << 30
|
||||||
|
sizeStr = strings.TrimSpace(sizeStr[:lastChar-1])
|
||||||
|
default:
|
||||||
|
multiplier = 1
|
||||||
|
sizeStr = strings.TrimSpace(sizeStr[:lastChar])
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
size := cast.ToInt(sizeStr)
|
||||||
|
if size < 0 {
|
||||||
|
size = 0
|
||||||
|
}
|
||||||
|
|
||||||
|
return safeMul(uint(size), multiplier)
|
||||||
|
}
|
||||||
|
|
6
viper.go
6
viper.go
|
@ -335,6 +335,12 @@ func (v *viper) GetStringMapString(key string) map[string]string {
|
||||||
return cast.ToStringMapString(v.Get(key))
|
return cast.ToStringMapString(v.Get(key))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func GetSizeInBytes(key string) uint { return v.GetSizeInBytes(key) }
|
||||||
|
func (v *viper) GetSizeInBytes(key string) uint {
|
||||||
|
sizeStr := cast.ToString(v.Get(key))
|
||||||
|
return parseSizeInBytes(sizeStr)
|
||||||
|
}
|
||||||
|
|
||||||
// Takes a single key and marshals it into a Struct
|
// Takes a single key and marshals it into a Struct
|
||||||
func MarshalKey(key string, rawVal interface{}) error { return v.MarshalKey(key, rawVal) }
|
func MarshalKey(key string, rawVal interface{}) error { return v.MarshalKey(key, rawVal) }
|
||||||
func (v *viper) MarshalKey(key string, rawVal interface{}) error {
|
func (v *viper) MarshalKey(key string, rawVal interface{}) error {
|
||||||
|
|
|
@ -376,3 +376,20 @@ func TestBoundCaseSensitivity(t *testing.T) {
|
||||||
assert.Equal(t, "green", Get("eyes"))
|
assert.Equal(t, "green", Get("eyes"))
|
||||||
|
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestSizeInBytes(t *testing.T) {
|
||||||
|
input := map[string]uint{
|
||||||
|
"": 0,
|
||||||
|
"b": 0,
|
||||||
|
"12 bytes": 0,
|
||||||
|
"200000000000gb": 0,
|
||||||
|
"12 b": 12,
|
||||||
|
"43 MB": 43 * (1 << 20),
|
||||||
|
"10mb": 10 * (1 << 20),
|
||||||
|
"1gb": 1 << 30,
|
||||||
|
}
|
||||||
|
|
||||||
|
for str, expected := range input {
|
||||||
|
assert.Equal(t, expected, parseSizeInBytes(str), str)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
Loading…
Reference in New Issue