flagext: add intstring type

CHANGELOG.md

@@ -98,6 +98,7 @@
 * [FEATURE] Add methods `Increment`, `FlushAll`, `CompareAndSwap`, `Touch` to `cache.MemcachedClient` #477
 * [FEATURE] Add `concurrency.ForEachJobMergeResults()` utility function. #486
 * [FEATURE] Add `ring.DoMultiUntilQuorumWithoutSuccessfulContextCancellation()`. #495
+* [FEATURE] Add `flagext.IntString` type. #616
 * [ENHANCEMENT] Add ability to log all source hosts from http header instead of only the first one. #444
 * [ENHANCEMENT] Add configuration to customize backoff for the gRPC clients.
 * [ENHANCEMENT] Use `SecretReader` interface to fetch secrets when configuring TLS. #274

flagext/intstring.go

@@ -0,0 +1,77 @@
+package flagext
+
+import (
+	"encoding/json"
+	"fmt"
+	"strconv"
+)
+
+// IntString is a type that can be unmarshaled from a string or an integer.
+type IntString int
+
+// UnmarshalYAML implements yaml.Unmarshaler.
+func (is *IntString) UnmarshalYAML(unmarshal func(interface{}) error) error {
+	var i int
+
+	if err := unmarshal(&i); err != nil {
+		var s string
+		if err = unmarshal(&s); err != nil {
+			return fmt.Errorf("IntString unmarshal error: %v", err)
+		}
+
+		if i, err = strconv.Atoi(s); err != nil {
+			return fmt.Errorf("IntString atoi error: %v", err)
+		}
+	}
+
+	*is = IntString(i)
+	return nil
+}
+
+// UnmarshalJSON implements json.Unmarshaler.
+func (is *IntString) UnmarshalJSON(data []byte) error {
+	if len(data) == 0 {
+		return nil
+	}
+
+	var i int
+
+	if err := json.Unmarshal(data, &i); err != nil {
+		var s string
+		if err = json.Unmarshal(data, &s); err != nil {
+			return fmt.Errorf("IntString unmarshal error: %v", err)
+		}
+
+		if i, err = strconv.Atoi(s); err != nil {
+			return fmt.Errorf("IntString atoi error: %v", err)
+		}
+	}
+
+	*is = IntString(i)
+	return nil
+}
+
+// String implements flag.Value.
+func (is *IntString) String() string {
+	var i int
+	if is != nil {
+		i = int(*is)
+	}
+	return fmt.Sprintf("%d", i)
+}
+
+// Set implements flag.Value.
+func (is *IntString) Set(val string) error {
+	if val == "" {
+		*is = 0
+		return nil
+	}
+
+	i, err := strconv.Atoi(val)
+	if err != nil {
+		return fmt.Errorf("IntString set error: %v", err)
+	}
+
+	*is = IntString(i)
+	return nil
+}

flagext/intstring_test.go

@@ -0,0 +1,156 @@
+package flagext
+
+import (
+	"fmt"
+	"strconv"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"gopkg.in/yaml.v2"
+)
+
+func Test_IntString_String(t *testing.T) {
+	for _, tcase := range []struct {
+		input    int
+		expected string
+	}{
+		{
+			input:    1024,
+			expected: "1024",
+		},
+		{
+			input:    2048000,
+			expected: "2048000",
+		},
+		{
+			input:    0,
+			expected: "0",
+		},
+	} {
+		t.Run(fmt.Sprintf("when input is %d", tcase.input), func(t *testing.T) {
+			b := IntString(tcase.input)
+			require.Equal(t, tcase.expected, b.String())
+		})
+	}
+}
+
+func Test_IntString_Set(t *testing.T) {
+	for _, tcase := range []struct {
+		input       string
+		expected    int
+		expectedErr bool
+	}{
+		{
+			input:    "2048000",
+			expected: 2048000,
+		},
+		{
+			input:    "512",
+			expected: 512,
+		},
+		{
+			input:    "40960000",
+			expected: 40960000,
+		},
+		{
+			expected: 0,
+		},
+		{
+			input:       "invalid",
+			expectedErr: true,
+		},
+	} {
+		t.Run(fmt.Sprintf("when input is %s", tcase.input), func(t *testing.T) {
+			var b IntString
+			require.True(t, (b.Set(tcase.input) != nil) == tcase.expectedErr)
+			require.Equal(t, b, IntString(tcase.expected))
+		})
+	}
+}
+
+func Test_IntString_UnmarshalYAML(t *testing.T) {
+	for _, tcase := range []struct {
+		input       []byte
+		expected    int
+		expectedErr bool
+	}{
+		{
+			input:    []byte(strconv.Itoa(0)),
+			expected: 0,
+		},
+		{
+			input:    []byte(strconv.Itoa(2048000)),
+			expected: 2048000,
+		},
+		{
+			input:    []byte("2048000"),
+			expected: 2048000,
+		},
+		{
+			input:    []byte(strconv.Itoa(40960000)),
+			expected: 40960000,
+		},
+		{
+			input:    []byte("40960000"),
+			expected: 40960000,
+		},
+		{
+			input:    []byte(""),
+			expected: 0,
+		},
+		{
+			input:       []byte("invalid"),
+			expectedErr: true,
+		},
+	} {
+		t.Run(fmt.Sprintf("when input is %s", tcase.input), func(t *testing.T) {
+			var b IntString
+			require.True(t, (yaml.Unmarshal(tcase.input, &b) != nil) == tcase.expectedErr)
+			require.Equal(t, IntString(tcase.expected), b)
+		})
+	}
+}
+
+func Test_IntString_UnmarshalJSON(t *testing.T) {
+	for _, tcase := range []struct {
+		input       []byte
+		expected    int
+		expectedErr bool
+	}{
+		{
+			input:    []byte(strconv.Itoa(0)),
+			expected: 0,
+		},
+		{
+			input:    []byte(strconv.Itoa(2048000)),
+			expected: 2048000,
+		},
+		{
+			input:    []byte("2048000"),
+			expected: 2048000,
+		},
+		{
+			input:    []byte(strconv.Itoa(40960000)),
+			expected: 40960000,
+		},
+		{
+			input:    []byte("40960000"),
+			expected: 40960000,
+		},
+		{
+			input:    []byte(""),
+			expected: 0,
+		},
+		{
+			input:       []byte("invalid"),
+			expectedErr: true,
+		},
+	} {
+		t.Run(fmt.Sprintf("when input is %s", tcase.input), func(t *testing.T) {
+			var b IntString
+
+			require.True(t, (b.UnmarshalJSON(tcase.input) != nil) == tcase.expectedErr)
+			require.Equal(t, IntString(tcase.expected), b)
+		})
+	}
+}