Merge pull request #395 from thockin/cleanups

IntOrString for use in JSON/YAML

pkg/util/util.go

@@ -64,3 +64,66 @@ func MakeJSONString(o interface{}) string {
 	data, _ := json.Marshal(o)
 	return string(data)
 }
+
+// IntOrString is a type that can hold an int or a string.  When used in
+// JSON or YAML marshalling and unmarshalling, it produces or consumes the
+// inner type.  This allows you to have, for example, a JSON field that can
+// accept a name or number.
+type IntOrString struct {
+	Kind   IntstrKind
+	IntVal int
+	StrVal string
+}
+
+type IntstrKind int
+
+const (
+	IntstrInt IntstrKind = iota
+	IntstrString
+)
+
+func (intstr *IntOrString) SetYAML(tag string, value interface{}) bool {
+	if intVal, ok := value.(int); ok {
+		intstr.Kind = IntstrInt
+		intstr.IntVal = intVal
+		return true
+	}
+	if strVal, ok := value.(string); ok {
+		intstr.Kind = IntstrString
+		intstr.StrVal = strVal
+		return true
+	}
+	return false
+}
+
+func (intstr IntOrString) GetYAML() (tag string, value interface{}) {
+	switch intstr.Kind {
+	case IntstrInt:
+		value = intstr.IntVal
+	case IntstrString:
+		value = intstr.StrVal
+	default:
+		panic("impossible IntOrString.Kind")
+	}
+	return
+}
+
+func (intstr *IntOrString) UnmarshalJSON(value []byte) error {
+	if value[0] == '"' {
+		intstr.Kind = IntstrString
+		return json.Unmarshal(value, &intstr.StrVal)
+	}
+	intstr.Kind = IntstrInt
+	return json.Unmarshal(value, &intstr.IntVal)
+}
+
+func (intstr IntOrString) MarshalJSON() ([]byte, error) {
+	switch intstr.Kind {
+	case IntstrInt:
+		return json.Marshal(intstr.IntVal)
+	case IntstrString:
+		return json.Marshal(intstr.StrVal)
+	default:
+		panic("impossible IntOrString.Kind")
+	}
+}

pkg/util/util_test.go

@@ -19,6 +19,8 @@ package util
 import (
 	"encoding/json"
 	"testing"
+
+	"gopkg.in/v1/yaml"
 )
 
 type FakeJSONBase struct {
@@ -65,3 +67,127 @@ func TestHandleCrash(t *testing.T) {
 		t.Errorf("Expected %d iterations, found %d", expect, count)
 	}
 }
+
+type IntOrStringHolder struct {
+	IOrS IntOrString `json:"val" yaml:"val"`
+}
+
+func TestIntOrStringUnmarshalYAML(t *testing.T) {
+	{
+		yaml_code_int := "val: 123\n"
+
+		var result IntOrStringHolder
+		if err := yaml.Unmarshal([]byte(yaml_code_int), &result); err != nil {
+			t.Errorf("Failed to unmarshal: %v", err)
+		}
+		if result.IOrS.Kind != IntstrInt || result.IOrS.IntVal != 123 {
+			t.Errorf("Failed to unmarshal int-typed IntOrString: %v", result)
+		}
+	}
+
+	{
+		yaml_code_str := "val: \"123\"\n"
+
+		var result IntOrStringHolder
+		if err := yaml.Unmarshal([]byte(yaml_code_str), &result); err != nil {
+			t.Errorf("Failed to unmarshal: %v", err)
+		}
+		if result.IOrS.Kind != IntstrString || result.IOrS.StrVal != "123" {
+			t.Errorf("Failed to unmarshal string-typed IntOrString: %v", result)
+		}
+	}
+}
+
+func TestIntOrStringMarshalYAML(t *testing.T) {
+	{
+		input := IntOrStringHolder{
+			IOrS: IntOrString{
+				Kind:   IntstrInt,
+				IntVal: 123,
+			},
+		}
+		result, err := yaml.Marshal(&input)
+		if err != nil {
+			t.Errorf("Failed to marshal: %v", err)
+		}
+		if string(result) != "val: 123\n" {
+			t.Errorf("Failed to marshal int-typed IntOrString: %q", string(result))
+		}
+	}
+
+	{
+		input := IntOrStringHolder{
+			IOrS: IntOrString{
+				Kind:   IntstrString,
+				StrVal: "123",
+			},
+		}
+		result, err := yaml.Marshal(&input)
+		if err != nil {
+			t.Errorf("Failed to marshal: %v", err)
+		}
+		if string(result) != "val: \"123\"\n" {
+			t.Errorf("Failed to marshal string-typed IntOrString: %q", string(result))
+		}
+	}
+}
+
+func TestIntOrStringUnmarshalJSON(t *testing.T) {
+	{
+		json_code_int := "{\"val\": 123}"
+
+		var result IntOrStringHolder
+		if err := json.Unmarshal([]byte(json_code_int), &result); err != nil {
+			t.Errorf("Failed to unmarshal: %v", err)
+		}
+		if result.IOrS.Kind != IntstrInt || result.IOrS.IntVal != 123 {
+			t.Errorf("Failed to unmarshal int-typed IntOrString: %v", result)
+		}
+	}
+
+	{
+		json_code_str := "{\"val\": \"123\"}"
+
+		var result IntOrStringHolder
+		if err := json.Unmarshal([]byte(json_code_str), &result); err != nil {
+			t.Errorf("Failed to unmarshal: %v", err)
+		}
+		if result.IOrS.Kind != IntstrString || result.IOrS.StrVal != "123" {
+			t.Errorf("Failed to unmarshal string-typed IntOrString: %v", result)
+		}
+	}
+}
+
+func TestIntOrStringMarshalJSON(t *testing.T) {
+	{
+		input := IntOrStringHolder{
+			IOrS: IntOrString{
+				Kind:   IntstrInt,
+				IntVal: 123,
+			},
+		}
+		result, err := json.Marshal(&input)
+		if err != nil {
+			t.Errorf("Failed to marshal: %v", err)
+		}
+		if string(result) != "{\"val\":123}" {
+			t.Errorf("Failed to marshal int-typed IntOrString: %q", string(result))
+		}
+	}
+
+	{
+		input := IntOrStringHolder{
+			IOrS: IntOrString{
+				Kind:   IntstrString,
+				StrVal: "123",
+			},
+		}
+		result, err := json.Marshal(&input)
+		if err != nil {
+			t.Errorf("Failed to marshal: %v", err)
+		}
+		if string(result) != "{\"val\":\"123\"}" {
+			t.Errorf("Failed to marshal string-typed IntOrString: %q", string(result))
+		}
+	}
+}