diff --git a/.travis.yml b/.travis.yml index 00d04cb9..89e77762 100644 --- a/.travis.yml +++ b/.travis.yml @@ -14,6 +14,7 @@ matrix: install: - go get golang.org/x/lint/golint + - go get github.com/docker/go-units - export PATH=$GOPATH/bin:$PATH - go install ./... diff --git a/byte_size.go b/byte_size.go new file mode 100644 index 00000000..cfcff3c4 --- /dev/null +++ b/byte_size.go @@ -0,0 +1,108 @@ +package pflag + +import ( + "github.com/docker/go-units" +) + +const byteSizeFlagType = "byte-size" + +// byteSizeValue used to pass byte sizes to a go-flags CLI +type byteSizeValue uint64 + +func newByteSizeValue(val uint64, p *uint64) *byteSizeValue { + *p = val + return (*byteSizeValue)(p) +} + +// MarshalFlag implements go-flags Marshaller interface +func (b byteSizeValue) MarshalFlag() (string, error) { + return units.HumanSize(float64(b)), nil +} + +// UnmarshalFlag implements go-flags Unmarshaller interface +func (b *byteSizeValue) UnmarshalFlag(value string) error { + sz, err := units.FromHumanSize(value) + if err != nil { + return err + } + *b = byteSizeValue(uint64(sz)) + return nil +} + +// String method for a bytesize (pflag value and stringer interface) +func (b byteSizeValue) String() string { + return units.HumanSize(float64(b)) +} + +// Set the value of this bytesize (pflag value interfaces) +func (b *byteSizeValue) Set(value string) error { + return b.UnmarshalFlag(value) +} + +// Type returns the type of the pflag value (pflag value interface) +func (b *byteSizeValue) Type() string { + return byteSizeFlagType +} + +func byteSizeConv(sval string) (interface{}, error) { + var b byteSizeValue + err := b.UnmarshalFlag(sval) + return uint64(b), err +} + +// GetByteSize return the ByteSize value of a flag with the given name +func (f *FlagSet) GetByteSize(name string) (uint64, error) { + val, err := f.getFlagType(name, byteSizeFlagType, byteSizeConv) + if err != nil { + return 0, err + } + return val.(uint64), nil +} + +// ByteSizeVar defines an uint64 flag with specified name, default value, and usage string. +// The argument p pouint64s to an uint64 variable in which to store the value of the flag. +func (f *FlagSet) ByteSizeVar(p *uint64, name string, value uint64, usage string) { + f.VarP(newByteSizeValue(value, p), name, "", usage) +} + +// ByteSizeVarP is like ByteSizeVar, but accepts a shorthand letter that can be used after a single dash. +func (f *FlagSet) ByteSizeVarP(p *uint64, name, shorthand string, value uint64, usage string) { + f.VarP(newByteSizeValue(value, p), name, shorthand, usage) +} + +// ByteSizeVar defines an uint64 flag with specified name, default value, and usage string. +// The argument p pouint64s to an uint64 variable in which to store the value of the flag. +func ByteSizeVar(p *uint64, name string, value uint64, usage string) { + CommandLine.VarP(newByteSizeValue(value, p), name, "", usage) +} + +// ByteSizeVarP is like ByteSizeVar, but accepts a shorthand letter that can be used after a single dash. +func ByteSizeVarP(p *uint64, name, shorthand string, value uint64, usage string) { + CommandLine.VarP(newByteSizeValue(value, p), name, shorthand, usage) +} + +// ByteSize defines an uint64 flag with specified name, default value, and usage string. +// The return value is the address of an uint64 variable that stores the value of the flag. +func (f *FlagSet) ByteSize(name string, value uint64, usage string) *uint64 { + p := new(uint64) + f.ByteSizeVarP(p, name, "", value, usage) + return p +} + +// ByteSizeP is like ByteSize, but accepts a shorthand letter that can be used after a single dash. +func (f *FlagSet) ByteSizeP(name, shorthand string, value uint64, usage string) *uint64 { + p := new(uint64) + f.ByteSizeVarP(p, name, shorthand, value, usage) + return p +} + +// ByteSize defines an uint64 flag with specified name, default value, and usage string. +// The return value is the address of an uint64 variable that stores the value of the flag. +func ByteSize(name string, value uint64, usage string) *uint64 { + return CommandLine.ByteSizeP(name, "", value, usage) +} + +// ByteSizeP is like ByteSize, but accepts a shorthand letter that can be used after a single dash. +func ByteSizeP(name, shorthand string, value uint64, usage string) *uint64 { + return CommandLine.ByteSizeP(name, shorthand, value, usage) +} diff --git a/byte_size_test.go b/byte_size_test.go new file mode 100644 index 00000000..2402dddd --- /dev/null +++ b/byte_size_test.go @@ -0,0 +1,117 @@ +package pflag + +import ( + "fmt" + "os" + "testing" + + "github.com/docker/go-units" +) + +func TestMarshalByteSize(t *testing.T) { + v, err := byteSizeValue(1024).MarshalFlag() + if err != nil { + t.Errorf("expected success, got %q", err) + } + expected := "1.024kB" + if v != expected { + t.Errorf("expected value to be %q, got %q", expected, v) + } +} + +func TestStringByteSize(t *testing.T) { + v := byteSizeValue(2048).String() + expected := "2.048kB" + if v != expected { + t.Errorf("expected value to be %q, got %q", expected, v) + } +} + +func TestUnmarshalByteSize(t *testing.T) { + var b byteSizeValue + err := b.UnmarshalFlag("notASize") + if err == nil { + t.Errorf("expected failure, got nil") + } + + err = b.UnmarshalFlag("1MB") + if err != nil { + t.Errorf("expected success, got %q", err) + } + expected := byteSizeValue(1000000) + if b != expected { + t.Errorf("expected value to be %d, got %d", expected, b) + } +} + +func TestSetByteSize(t *testing.T) { + var b byteSizeValue + err := b.Set("notASize") + if err == nil { + t.Errorf("expected failure, got nil") + } + + err = b.Set("2MB") + if err != nil { + t.Errorf("expected success, got %q", err) + } + expected := byteSizeValue(2000000) + if b != expected { + t.Errorf("expected value to be %d, got %d", expected, b) + } +} + +func TestTypeByteSize(t *testing.T) { + var b byteSizeValue + v := b.Type() + expected := "byte-size" + if v != expected { + t.Errorf("expected value to be %q, got %q", expected, v) + } +} + +func setUpByteSize(value *uint64) *FlagSet { + f := NewFlagSet("test", ContinueOnError) + f.ByteSizeVar(value, "size", 1*units.MiB, "Size") + return f +} + +func TestByteSize(t *testing.T) { + testCases := []struct { + input string + success bool + expected uint64 + }{ + {"1KB", true, 1000}, + {"1MB", true, 1000000}, + {"1kb", true, 1000}, + {"zzz", false, 0}, + } + + devnull, _ := os.Open(os.DevNull) + os.Stderr = devnull + for i := range testCases { + var addr uint64 + f := setUpByteSize(&addr) + + tc := &testCases[i] + + arg := fmt.Sprintf("--size=%s", tc.input) + err := f.Parse([]string{arg}) + if err != nil && tc.success == true { + t.Errorf("expected success, got %q", err) + continue + } else if err == nil && tc.success == false { + t.Errorf("expected failure") + continue + } else if tc.success { + size, err := f.GetByteSize("size") + if err != nil { + t.Errorf("Got error trying to fetch the IP flag: %v", err) + } + if size != tc.expected { + t.Errorf("for input %q, expected %d, got %d", tc.input, tc.expected, size) + } + } + } +} diff --git a/go.mod b/go.mod index b2287eec..81854bf0 100644 --- a/go.mod +++ b/go.mod @@ -1,3 +1,5 @@ module github.com/spf13/pflag go 1.12 + +require github.com/docker/go-units v0.4.0 diff --git a/go.sum b/go.sum index e69de29b..e1d7438f 100644 --- a/go.sum +++ b/go.sum @@ -0,0 +1,2 @@ +github.com/docker/go-units v0.4.0 h1:3uh0PgVws3nIA0Q+MwDC8yjEPf9zjRfZZWXZYDct3Tw= +github.com/docker/go-units v0.4.0/go.mod h1:fgPhTUdO+D/Jk86RDLlptpiXQzgHJF7gydDDbaIK4Dk=