From 608f24a6e3d43e596416db046cb24df20da65405 Mon Sep 17 00:00:00 2001
From: iwilltry42 <iwilltry42@gmail.com>
Date: Mon, 18 Jan 2021 08:52:47 +0100
Subject: [PATCH] completion/zsh: explicitly source zsh completion

Thanks for the hint @konradmalik, fixes #390
---
 cmd/root.go | 16 ++++++++++++----
 1 file changed, 12 insertions(+), 4 deletions(-)

diff --git a/cmd/root.go b/cmd/root.go
index cd9cdf50..bdd84bcc 100644
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -184,8 +184,16 @@ func generateFishCompletion(writer io.Writer) error {
 
 // Completion
 var completionFunctions = map[string]func(io.Writer) error{
-	"bash":       rootCmd.GenBashCompletion,
-	"zsh":        rootCmd.GenZshCompletion,
+	"bash": rootCmd.GenBashCompletion,
+	"zsh": func(writer io.Writer) error {
+		if err := rootCmd.GenZshCompletion(writer); err != nil {
+			return err
+		}
+
+		fmt.Fprintf(writer, "\n# source completion file\ncompdef _k3d k3d\n")
+
+		return nil
+	},
 	"psh":        rootCmd.GenPowerShellCompletion,
 	"powershell": rootCmd.GenPowerShellCompletion,
 	"fish":       generateFishCompletion,
@@ -200,8 +208,8 @@ func NewCmdCompletion() *cobra.Command {
 		Long:  `Generate completion scripts for [bash, zsh, fish, powershell | psh]`,
 		Args:  cobra.ExactArgs(1), // TODO: NewCmdCompletion: add support for 0 args = auto detection
 		Run: func(cmd *cobra.Command, args []string) {
-			if f, ok := completionFunctions[args[0]]; ok {
-				if err := f(os.Stdout); err != nil {
+			if completionFunc, ok := completionFunctions[args[0]]; ok {
+				if err := completionFunc(os.Stdout); err != nil {
 					log.Fatalf("Failed to generate completion script for shell '%s'", args[0])
 				}
 				return