From 8092588fb88b8c409e017be2350b2f40f0c17865 Mon Sep 17 00:00:00 2001
From: "Dr. Stefan Schimanski" <stefan.schimanski@gmail.com>
Date: Sun, 20 Mar 2016 21:05:18 +0100
Subject: [PATCH] Add custom bash flag completion

---
 bash_completions.go      | 27 +++++++++++++++++++++++++++
 bash_completions.md      | 32 ++++++++++++++++++++++++++++++++
 bash_completions_test.go |  7 +++++++
 3 files changed, 66 insertions(+)

diff --git a/bash_completions.go b/bash_completions.go
index 2e46356..61fd88b 100644
--- a/bash_completions.go
+++ b/bash_completions.go
@@ -12,6 +12,7 @@ import (
 
 const (
 	BashCompFilenameExt     = "cobra_annotation_bash_completion_filename_extentions"
+	BashCompCustom          = "cobra_annotation_bash_completion_custom"
 	BashCompOneRequiredFlag = "cobra_annotation_bash_completion_one_required_flag"
 	BashCompSubdirsInDir    = "cobra_annotation_bash_completion_subdirs_in_dir"
 )
@@ -318,6 +319,20 @@ func writeFlagHandler(name string, annotations map[string][]string, w io.Writer)
 			if err != nil {
 				return err
 			}
+		case BashCompCustom:
+			_, err := fmt.Fprintf(w, "    flags_with_completion+=(%q)\n", name)
+			if err != nil {
+				return err
+			}
+			if len(value) > 0 {
+				handlers := strings.Join(value, "; ")
+				_, err = fmt.Fprintf(w, "    flags_completion+=(%q)\n", handlers)
+			} else {
+				_, err = fmt.Fprintf(w, "    flags_completion+=(:)\n")
+			}
+			if err != nil {
+				return err
+			}
 		case BashCompSubdirsInDir:
 			_, err := fmt.Fprintf(w, "    flags_with_completion+=(%q)\n", name)
 
@@ -538,6 +553,12 @@ func (cmd *Command) MarkFlagFilename(name string, extensions ...string) error {
 	return MarkFlagFilename(cmd.Flags(), name, extensions...)
 }
 
+// MarkFlagCustom adds the BashCompCustom annotation to the named flag, if it exists.
+// Generated bash autocompletion will call the bash function f for the flag.
+func (cmd *Command) MarkFlagCustom(name string, f string) error {
+	return MarkFlagCustom(cmd.Flags(), name, f)
+}
+
 // MarkPersistentFlagFilename adds the BashCompFilenameExt annotation to the named persistent flag, if it exists.
 // Generated bash autocompletion will select filenames for the flag, limiting to named extensions if provided.
 func (cmd *Command) MarkPersistentFlagFilename(name string, extensions ...string) error {
@@ -549,3 +570,9 @@ func (cmd *Command) MarkPersistentFlagFilename(name string, extensions ...string
 func MarkFlagFilename(flags *pflag.FlagSet, name string, extensions ...string) error {
 	return flags.SetAnnotation(name, BashCompFilenameExt, extensions)
 }
+
+// MarkFlagCustom adds the BashCompCustom annotation to the named flag in the flag set, if it exists.
+// Generated bash autocompletion will call the bash function f for the flag.
+func MarkFlagCustom(flags *pflag.FlagSet, name string, f string) error {
+	return flags.SetAnnotation(name, BashCompCustom, []string{f})
+}
diff --git a/bash_completions.md b/bash_completions.md
index 204704e..691b1b7 100644
--- a/bash_completions.md
+++ b/bash_completions.md
@@ -147,3 +147,35 @@ hello.yml                     test.json
 ```
 
 So while there are many other files in the CWD it only shows me subdirs and those with valid extensions.
+
+# Specifiy custom flag completion
+
+Similar to the filename completion and filtering usingn cobra.BashCompFilenameExt, you can specifiy
+a custom flag completion function with cobra.BashCompCustom:
+
+```go
+	annotation := make(map[string][]string)
+	annotation[cobra.BashCompFilenameExt] = []string{"__kubectl_get_namespaces"}
+
+	flag := &pflag.Flag{
+		Name:        "namespace",
+		Usage:       usage,
+		Annotations: annotation,
+	}
+	cmd.Flags().AddFlag(flag)
+```
+
+In addition add the `__handle_namespace_flag` implementation in the `BashCompletionFunction`
+value, e.g.:
+
+```bash
+__kubectl_get_namespaces()
+{
+    local template
+    template="{{ range .items  }}{{ .metadata.name }} {{ end }}"
+    local kubectl_out
+    if kubectl_out=$(kubectl get -o template --template="${template}" namespace 2>/dev/null); then
+        COMPREPLY=( $( compgen -W "${kubectl_out}[*]" -- "$cur" ) )
+    fi
+}
+```
diff --git a/bash_completions_test.go b/bash_completions_test.go
index 86f3d01..cf2661c 100644
--- a/bash_completions_test.go
+++ b/bash_completions_test.go
@@ -62,6 +62,11 @@ func TestBashCompletions(t *testing.T) {
 	c.Flags().StringVar(&flagvalExt, "filename-ext", "", "Enter a filename (extension limited)")
 	c.MarkFlagFilename("filename-ext")
 
+	// filename extensions
+	var flagvalCustom string
+	c.Flags().StringVar(&flagvalCustom, "custom", "", "Enter a filename (extension limited)")
+	c.MarkFlagCustom("custom", "__complete_custom")
+
 	// subdirectories in a given directory
 	var flagvalTheme string
 	c.Flags().StringVar(&flagvalTheme, "theme", "", "theme to use (located in /themes/THEMENAME/)")
@@ -88,6 +93,8 @@ func TestBashCompletions(t *testing.T) {
 	check(t, str, `flags_completion+=("_filedir")`)
 	// check for filename extension flags
 	check(t, str, `flags_completion+=("__handle_filename_extension_flag json|yaml|yml")`)
+	// check for custom flags
+	check(t, str, `flags_completion+=("__complete_custom")`)
 	// check for subdirs_in_dir flags
 	check(t, str, `flags_completion+=("__handle_subdirs_in_dir_flag themes")`)