From f1557d104d9e0660867ca45a2ce76d28905d902b Mon Sep 17 00:00:00 2001
From: Adphi <philippe.adrien.nousse@gmail.com>
Date: Mon, 12 Sep 2022 10:01:04 +0200
Subject: [PATCH] fix zsh completion not working with `source` README.md: add
 shell completion installation instructions

Signed-off-by: Adphi <philippe.adrien.nousse@gmail.com>
---
 README.md        | 17 +++++++++++++++++
 cmd/d2vm/main.go |  6 ++++++
 2 files changed, 23 insertions(+)

diff --git a/README.md b/README.md
index 48b1c15..048518d 100644
--- a/README.md
+++ b/README.md
@@ -75,6 +75,23 @@ which d2vm
 /go/bin/d2vm
 ```
 
+### Generate shell completion
+
+The *d2vm* program supports shell completion for *bash*, *zsh* and *fish*.
+
+It can be enabled by running the following command:
+
+```bash
+source <(d2vm completion $(basename $SHELL))
+```
+
+Or you can install the completion file in the shell completion directory by following the instructions:
+
+```bash
+d2vm completion $(basename $SHELL) --help
+```
+
+
 ### Converting an existing Docker Image to VM image:
 
 ```bash
diff --git a/cmd/d2vm/main.go b/cmd/d2vm/main.go
index 9eae0d6..5950808 100644
--- a/cmd/d2vm/main.go
+++ b/cmd/d2vm/main.go
@@ -56,6 +56,12 @@ var (
 				logrus.SetLevel(logrus.TraceLevel)
 			}
 			exec.SetDebug(verbose)
+
+			// make the zsh completion work when sourced with `source <(d2vm completion zsh)`
+			if cmd.Name() == "zsh" && cmd.Parent() != nil && cmd.Parent().Name() == "completion" {
+				zshHead := fmt.Sprintf("#compdef %[1]s\ncompdef _%[1]s %[1]s\n", cmd.Root().Name())
+				cmd.OutOrStdout().Write([]byte(zshHead))
+			}
 		},
 	}
 )