flag: allow CommandLine's Usage function to be set

Fixes issue 7779.

LGTM=rsc
R=golang-codereviews, rsc
CC=golang-codereviews
https://codereview.appspot.com/147210043

2 files changed, 19 insertions, 6 deletions

diff --git a/src/flag/flag.go b/src/flag/flag.go
index de2d91f8b..323e452a8 100644
--- a/src/flag/flag.go
+++ b/src/flag/flag.go
@@ -406,6 +406,7 @@ func defaultUsage(f *FlagSet) {
 // for how to write your own usage function.
 
 // Usage prints to standard error a usage message documenting all defined command-line flags.
+// It is called when an error occurs while parsing flags.
 // The function is a variable that may be changed to point to a custom function.
 var Usage = func() {
 	fmt.Fprintf(os.Stderr, "Usage of %s:\n", os.Args[0])
@@ -702,13 +703,15 @@ func (f *FlagSet) failf(format string, a ...interface{}) error {
 	return err
 }
 
-// usage calls the Usage method for the flag set, or the usage function if
-// the flag set is CommandLine.
+// usage calls the Usage method for the flag set if one is specified,
+// or the appropriate default usage function otherwise.
 func (f *FlagSet) usage() {
-	if f == CommandLine {
-		Usage()
-	} else if f.Usage == nil {
-		defaultUsage(f)
+	if f.Usage == nil {
+		if f == CommandLine {
+			Usage()
+		} else {
+			defaultUsage(f)
+		}
 	} else {
 		f.Usage()
 	}
diff --git a/src/flag/flag_test.go b/src/flag/flag_test.go
index 2c0387269..8c88c8c27 100644
--- a/src/flag/flag_test.go
+++ b/src/flag/flag_test.go
@@ -251,6 +251,16 @@ func TestUserDefined(t *testing.T) {
 	}
 }
 
+func TestUserDefinedForCommandLine(t *testing.T) {
+	const help = "HELP"
+	var result string
+	ResetForTesting(func() { result = help })
+	Usage()
+	if result != help {
+		t.Fatalf("got %q; expected %q", result, help)
+	}
+}
+
 // Declare a user-defined boolean flag type.
 type boolFlagVar struct {
 	count int