[#857] config: allow to parse sizes with suffix

The code is taken from viper as `cast` package that we use
doesn't have needed converter.

Signed-off-by: Evgenii Stratonikov <evgeniy@nspcc.ru>
This commit is contained in:
Evgenii Stratonikov 2021-10-07 19:15:48 +03:00 committed by Leonard Lyubich
parent b8ba677c85
commit a0abf10c8e
4 changed files with 100 additions and 0 deletions

View file

@ -1,7 +1,9 @@
package config package config
import ( import (
"strings"
"time" "time"
"unicode"
"github.com/spf13/cast" "github.com/spf13/cast"
) )
@ -144,3 +146,63 @@ func Int(c *Config, name string) int64 {
func IntSafe(c *Config, name string) int64 { func IntSafe(c *Config, name string) int64 {
return cast.ToInt64(c.Value(name)) return cast.ToInt64(c.Value(name))
} }
// SizeInBytesSafe reads configuration value
// from c by name and casts it to size in bytes (uint64).
//
// The suffix can be single-letter (b, k, m, g, t) or with
// an additional b at the end. Spaces between the number and the suffix
// are allowed. All multipliers are power of 2 (i.e. k is for kibi-byte).
//
// Returns 0 if a value can't be casted.
func SizeInBytesSafe(c *Config, name string) uint64 {
s := StringSafe(c, name)
return parseSizeInBytes(s)
}
// The following code is taken from https://github.com/spf13/viper/blob/master/util.go
// with minor corrections (allow to use both `k` and `kb` forms.
// Seems like viper allows to convert sizes but corresponding parser in `cast` package
// is missing.
func safeMul(a, b uint64) uint64 {
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) uint64 {
sizeStr = strings.TrimSpace(sizeStr)
lastChar := len(sizeStr) - 1
multiplier := uint64(1)
if lastChar > 0 {
if sizeStr[lastChar] == 'b' || sizeStr[lastChar] == 'B' {
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])
case 't':
multiplier = 1 << 40
sizeStr = strings.TrimSpace(sizeStr[:lastChar])
default:
multiplier = 1
sizeStr = strings.TrimSpace(sizeStr[:lastChar+1])
}
}
}
size := cast.ToUint64(sizeStr)
return safeMul(size, multiplier)
}

View file

@ -115,3 +115,22 @@ func TestNumbers(t *testing.T) {
require.Zero(t, config.UintSafe(c, incorrect)) require.Zero(t, config.UintSafe(c, incorrect))
}) })
} }
func TestSizeInBytes(t *testing.T) {
const (
kb = 1024
mb = 1024 * kb
gb = 1024 * mb
tb = 1024 * gb
)
configtest.ForEachFileType("test/config", func(c *config.Config) {
c = c.Sub("sizes")
require.EqualValues(t, kb, config.SizeInBytesSafe(c, "size_kb"))
require.EqualValues(t, 2*kb, config.SizeInBytesSafe(c, "size_kb_no_space"))
require.EqualValues(t, 12*mb, config.SizeInBytesSafe(c, "size_mb"))
require.EqualValues(t, 4*gb, config.SizeInBytesSafe(c, "size_gb"))
require.EqualValues(t, 5*tb, config.SizeInBytesSafe(c, "size_tb"))
require.EqualValues(t, 2048, config.SizeInBytesSafe(c, "size_bytes"))
require.EqualValues(t, 123456, config.SizeInBytesSafe(c, "size_bytes_no_suffix"))
})
}

View file

@ -46,5 +46,15 @@
"correct": true, "correct": true,
"correct_string": "true", "correct_string": "true",
"incorrect": "not true" "incorrect": "not true"
},
"sizes": {
"size_kb": "1 kb",
"size_kb_no_space": "2kb",
"size_mb": "12m",
"size_gb": "4g",
"size_tb": "5 TB",
"size_bytes": "2048b",
"size_bytes_no_suffix": 123456
} }
} }

View file

@ -40,3 +40,12 @@ bool:
correct: true correct: true
correct_string: "true" correct_string: "true"
incorrect: not true incorrect: not true
sizes:
size_kb: 1 kb
size_kb_no_space: 2kb
size_mb: 12m
size_gb: 4g
size_tb: 5 TB
size_bytes: 2048b
size_bytes_no_suffix: 123456