Add SetAnnotation helper

2 files changed, 45 insertions, 0 deletions

diff --git a/flag.go b/flag.go
index 0070b93..42463c6 100644
--- a/flag.go
+++ b/flag.go
@@ -295,6 +295,19 @@ func (f *FlagSet) Set(name, value string) error {
 	return nil
 }
 
+func (f *FlagSet) SetAnnotation(name, key string, values []string) error {
+	normalName := f.normalizeFlagName(name)
+	flag, ok := f.formal[normalName]
+	if !ok {
+		return fmt.Errorf("no such flag -%v", name)
+	}
+	if flag.Annotations == nil {
+		flag.Annotations = map[string][]string{}
+	}
+	flag.Annotations[key] = values
+	return nil
+}
+
 // Set sets the value of the named command-line flag.
 func Set(name, value string) error {
 	return CommandLine.Set(name, value)
diff --git a/flag_test.go b/flag_test.go
index d3c1714..40140d2 100644
--- a/flag_test.go
+++ b/flag_test.go
@@ -11,6 +11,7 @@ import (
 	"io/ioutil"
 	"net"
 	"os"
+	"reflect"
 	"sort"
 	"strings"
 	"testing"
@@ -107,6 +108,37 @@ func TestUsage(t *testing.T) {
 	}
 }
 
+func TestAnnotation(t *testing.T) {
+	f := NewFlagSet("shorthand", ContinueOnError)
+
+	if err := f.SetAnnotation("missing-flag", "key", nil); err == nil {
+		t.Errorf("Expected error setting annotation on non-existent flag")
+	}
+
+	f.StringP("stringa", "a", "", "string value")
+	if err := f.SetAnnotation("stringa", "key", nil); err != nil {
+		t.Errorf("Unexpected error setting new nil annotation: %v", err)
+	}
+	if annotation := f.Lookup("stringa").Annotations["key"]; annotation != nil {
+		t.Errorf("Unexpected annotation: %v", annotation)
+	}
+
+	f.StringP("stringb", "b", "", "string2 value")
+	if err := f.SetAnnotation("stringb", "key", []string{"value1"}); err != nil {
+		t.Errorf("Unexpected error setting new annotation: %v", err)
+	}
+	if annotation := f.Lookup("stringb").Annotations["key"]; !reflect.DeepEqual(annotation, []string{"value1"}) {
+		t.Errorf("Unexpected annotation: %v", annotation)
+	}
+
+	if err := f.SetAnnotation("stringb", "key", []string{"value2"}); err != nil {
+		t.Errorf("Unexpected error updating annotation: %v", err)
+	}
+	if annotation := f.Lookup("stringb").Annotations["key"]; !reflect.DeepEqual(annotation, []string{"value2"}) {
+		t.Errorf("Unexpected annotation: %v", annotation)
+	}
+}
+
 func testParse(f *FlagSet, t *testing.T) {
 	if f.Parsed() {
 		t.Error("f.Parse() = true before Parse")