Prevent changing order of flags after SetNormalizeFunc() (#121)

Fix #120

2 files changed, 10 insertions, 6 deletions

diff --git a/flag.go b/flag.go
index cb57359..3a2e255 100644
--- a/flag.go
+++ b/flag.go
@@ -200,13 +200,12 @@ func sortFlags(flags map[NormalizedName]*Flag) []*Flag {
 // "--getUrl" which may also be translated to "geturl" and everything will work.
 func (f *FlagSet) SetNormalizeFunc(n func(f *FlagSet, name string) NormalizedName) {
 	f.normalizeNameFunc = n
-	f.orderedFormal = f.orderedFormal[:0]
-	for k, v := range f.formal {
-		delete(f.formal, k)
-		nname := f.normalizeFlagName(string(k))
-		f.formal[nname] = v
-		f.orderedFormal = append(f.orderedFormal, v)
+	for k, v := range f.orderedFormal {
+		delete(f.formal, NormalizedName(v.Name))
+		nname := f.normalizeFlagName(v.Name)
 		v.Name = string(nname)
+		f.formal[nname] = v
+		f.orderedFormal[k] = v
 	}
 }
 
diff --git a/flag_test.go b/flag_test.go
index 55dd6c6..09249c8 100644
--- a/flag_test.go
+++ b/flag_test.go
@@ -1008,6 +1008,11 @@ func TestPrintDefaults(t *testing.T) {
 func TestVisitAllFlagOrder(t *testing.T) {
 	fs := NewFlagSet("TestVisitAllFlagOrder", ContinueOnError)
 	fs.SortFlags = false
+	// https://github.com/spf13/pflag/issues/120
+	fs.SetNormalizeFunc(func(f *FlagSet, name string) NormalizedName {
+		return NormalizedName(name)
+	})
+
 	names := []string{"C", "B", "A", "D"}
 	for _, name := range names {
 		fs.Bool(name, false, "")