Added support for the byte-size type

.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 ./...
 

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)
+}

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)
+			}
+		}
+	}
+}

go.mod

@@ -1,3 +1,5 @@
 module github.com/spf13/pflag
 
 go 1.12
+
+require github.com/docker/go-units v0.4.0

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=