From 5b6f88655f45a2091754d7a155f3d36c0f42f0cf Mon Sep 17 00:00:00 2001 From: Aleksei Besogonov Date: Thu, 27 Jun 2019 18:58:06 -0700 Subject: [PATCH] Added String-To-Int64 option parsing --- string_to_int64.go | 149 ++++++++++++++++++++++++++++++++++++++ string_to_int64_test.go | 156 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 305 insertions(+) create mode 100644 string_to_int64.go create mode 100644 string_to_int64_test.go diff --git a/string_to_int64.go b/string_to_int64.go new file mode 100644 index 00000000..a807a04a --- /dev/null +++ b/string_to_int64.go @@ -0,0 +1,149 @@ +package pflag + +import ( + "bytes" + "fmt" + "strconv" + "strings" +) + +// -- stringToInt64 Value +type stringToInt64Value struct { + value *map[string]int64 + changed bool +} + +func newStringToInt64Value(val map[string]int64, p *map[string]int64) *stringToInt64Value { + ssv := new(stringToInt64Value) + ssv.value = p + *ssv.value = val + return ssv +} + +// Format: a=1,b=2 +func (s *stringToInt64Value) Set(val string) error { + ss := strings.Split(val, ",") + out := make(map[string]int64, len(ss)) + for _, pair := range ss { + kv := strings.SplitN(pair, "=", 2) + if len(kv) != 2 { + return fmt.Errorf("%s must be formatted as key=value", pair) + } + var err error + out[kv[0]], err = strconv.ParseInt(kv[1], 10, 64) + if err != nil { + return err + } + } + if !s.changed { + *s.value = out + } else { + for k, v := range out { + (*s.value)[k] = v + } + } + s.changed = true + return nil +} + +func (s *stringToInt64Value) Type() string { + return "stringToInt64" +} + +func (s *stringToInt64Value) String() string { + var buf bytes.Buffer + i := 0 + for k, v := range *s.value { + if i > 0 { + buf.WriteRune(',') + } + buf.WriteString(k) + buf.WriteRune('=') + buf.WriteString(strconv.FormatInt(v, 10)) + i++ + } + return "[" + buf.String() + "]" +} + +func stringToInt64Conv(val string) (interface{}, error) { + val = strings.Trim(val, "[]") + // An empty string would cause an empty map + if len(val) == 0 { + return map[string]int64{}, nil + } + ss := strings.Split(val, ",") + out := make(map[string]int64, len(ss)) + for _, pair := range ss { + kv := strings.SplitN(pair, "=", 2) + if len(kv) != 2 { + return nil, fmt.Errorf("%s must be formatted as key=value", pair) + } + var err error + out[kv[0]], err = strconv.ParseInt(kv[1], 10, 64) + if err != nil { + return nil, err + } + } + return out, nil +} + +// GetStringToInt64 return the map[string]int64 value of a flag with the given name +func (f *FlagSet) GetStringToInt64(name string) (map[string]int64, error) { + val, err := f.getFlagType(name, "stringToInt64", stringToInt64Conv) + if err != nil { + return map[string]int64{}, err + } + return val.(map[string]int64), nil +} + +// StringToInt64Var defines a string flag with specified name, default value, and usage string. +// The argument p point64s to a map[string]int64 variable in which to store the values of the multiple flags. +// The value of each argument will not try to be separated by comma +func (f *FlagSet) StringToInt64Var(p *map[string]int64, name string, value map[string]int64, usage string) { + f.VarP(newStringToInt64Value(value, p), name, "", usage) +} + +// StringToInt64VarP is like StringToInt64Var, but accepts a shorthand letter that can be used after a single dash. +func (f *FlagSet) StringToInt64VarP(p *map[string]int64, name, shorthand string, value map[string]int64, usage string) { + f.VarP(newStringToInt64Value(value, p), name, shorthand, usage) +} + +// StringToInt64Var defines a string flag with specified name, default value, and usage string. +// The argument p point64s to a map[string]int64 variable in which to store the value of the flag. +// The value of each argument will not try to be separated by comma +func StringToInt64Var(p *map[string]int64, name string, value map[string]int64, usage string) { + CommandLine.VarP(newStringToInt64Value(value, p), name, "", usage) +} + +// StringToInt64VarP is like StringToInt64Var, but accepts a shorthand letter that can be used after a single dash. +func StringToInt64VarP(p *map[string]int64, name, shorthand string, value map[string]int64, usage string) { + CommandLine.VarP(newStringToInt64Value(value, p), name, shorthand, usage) +} + +// StringToInt64 defines a string flag with specified name, default value, and usage string. +// The return value is the address of a map[string]int64 variable that stores the value of the flag. +// The value of each argument will not try to be separated by comma +func (f *FlagSet) StringToInt64(name string, value map[string]int64, usage string) *map[string]int64 { + p := map[string]int64{} + f.StringToInt64VarP(&p, name, "", value, usage) + return &p +} + +// StringToInt64P is like StringToInt64, but accepts a shorthand letter that can be used after a single dash. +func (f *FlagSet) StringToInt64P(name, shorthand string, value map[string]int64, usage string) *map[string]int64 { + p := map[string]int64{} + f.StringToInt64VarP(&p, name, shorthand, value, usage) + return &p +} + +// StringToInt64 defines a string flag with specified name, default value, and usage string. +// The return value is the address of a map[string]int64 variable that stores the value of the flag. +// The value of each argument will not try to be separated by comma +func StringToInt64(name string, value map[string]int64, usage string) *map[string]int64 { + return CommandLine.StringToInt64P(name, "", value, usage) +} + +// StringToInt64P is like StringToInt64, but accepts a shorthand letter that can be used after a single dash. +func StringToInt64P(name, shorthand string, value map[string]int64, usage string) *map[string]int64 { + return CommandLine.StringToInt64P(name, shorthand, value, usage) +} diff --git a/string_to_int64_test.go b/string_to_int64_test.go new file mode 100644 index 00000000..2b3f2989 --- /dev/null +++ b/string_to_int64_test.go @@ -0,0 +1,156 @@ +// Copyright 2009 The Go Authors. All rights reserved. +// Use of ths2i source code s2i governed by a BSD-style +// license that can be found in the LICENSE file. + +package pflag + +import ( + "bytes" + "fmt" + "strconv" + "testing" +) + +func setUpS2I64FlagSet(s2ip *map[string]int64) *FlagSet { + f := NewFlagSet("test", ContinueOnError) + f.StringToInt64Var(s2ip, "s2i", map[string]int64{}, "Command separated ls2it!") + return f +} + +func setUpS2I64FlagSetWithDefault(s2ip *map[string]int64) *FlagSet { + f := NewFlagSet("test", ContinueOnError) + f.StringToInt64Var(s2ip, "s2i", map[string]int64{"a": 1, "b": 2}, "Command separated ls2it!") + return f +} + +func createS2I64Flag(vals map[string]int64) string { + var buf bytes.Buffer + i := 0 + for k, v := range vals { + if i > 0 { + buf.WriteRune(',') + } + buf.WriteString(k) + buf.WriteRune('=') + buf.WriteString(strconv.FormatInt(v, 10)) + i++ + } + return buf.String() +} + +func TestEmptyS2I64(t *testing.T) { + var s2i map[string]int64 + f := setUpS2I64FlagSet(&s2i) + err := f.Parse([]string{}) + if err != nil { + t.Fatal("expected no error; got", err) + } + + getS2I, err := f.GetStringToInt64("s2i") + if err != nil { + t.Fatal("got an error from GetStringToInt64():", err) + } + if len(getS2I) != 0 { + t.Fatalf("got s2i %v with len=%d but expected length=0", getS2I, len(getS2I)) + } +} + +func TestS2I64(t *testing.T) { + var s2i map[string]int64 + f := setUpS2I64FlagSet(&s2i) + + vals := map[string]int64{"a": 1, "b": 2, "d": 4, "c": 3} + arg := fmt.Sprintf("--s2i=%s", createS2I64Flag(vals)) + err := f.Parse([]string{arg}) + if err != nil { + t.Fatal("expected no error; got", err) + } + for k, v := range s2i { + if vals[k] != v { + t.Fatalf("expected s2i[%s] to be %d but got: %d", k, vals[k], v) + } + } + getS2I, err := f.GetStringToInt64("s2i") + if err != nil { + t.Fatalf("got error: %v", err) + } + for k, v := range getS2I { + if vals[k] != v { + t.Fatalf("expected s2i[%s] to be %d but got: %d from GetStringToInt64", k, vals[k], v) + } + } +} + +func TestS2I64Default(t *testing.T) { + var s2i map[string]int64 + f := setUpS2I64FlagSetWithDefault(&s2i) + + vals := map[string]int64{"a": 1, "b": 2} + + err := f.Parse([]string{}) + if err != nil { + t.Fatal("expected no error; got", err) + } + for k, v := range s2i { + if vals[k] != v { + t.Fatalf("expected s2i[%s] to be %d but got: %d", k, vals[k], v) + } + } + + getS2I, err := f.GetStringToInt64("s2i") + if err != nil { + t.Fatal("got an error from GetStringToInt64():", err) + } + for k, v := range getS2I { + if vals[k] != v { + t.Fatalf("expected s2i[%s] to be %d from GetStringToInt64 but got: %d", k, vals[k], v) + } + } +} + +func TestS2I64WithDefault(t *testing.T) { + var s2i map[string]int64 + f := setUpS2I64FlagSetWithDefault(&s2i) + + vals := map[string]int64{"a": 1, "b": 2} + arg := fmt.Sprintf("--s2i=%s", createS2I64Flag(vals)) + err := f.Parse([]string{arg}) + if err != nil { + t.Fatal("expected no error; got", err) + } + for k, v := range s2i { + if vals[k] != v { + t.Fatalf("expected s2i[%s] to be %d but got: %d", k, vals[k], v) + } + } + + getS2I, err := f.GetStringToInt64("s2i") + if err != nil { + t.Fatal("got an error from GetStringToInt64():", err) + } + for k, v := range getS2I { + if vals[k] != v { + t.Fatalf("expected s2i[%s] to be %d from GetStringToInt64 but got: %d", k, vals[k], v) + } + } +} + +func TestS2I64CalledTwice(t *testing.T) { + var s2i map[string]int64 + f := setUpS2I64FlagSet(&s2i) + + in := []string{"a=1,b=2", "b=3"} + expected := map[string]int64{"a": 1, "b": 3} + argfmt := "--s2i=%s" + arg1 := fmt.Sprintf(argfmt, in[0]) + arg2 := fmt.Sprintf(argfmt, in[1]) + err := f.Parse([]string{arg1, arg2}) + if err != nil { + t.Fatal("expected no error; got", err) + } + for i, v := range s2i { + if expected[i] != v { + t.Fatalf("expected s2i[%s] to be %d but got: %d", i, expected[i], v) + } + } +}