From f99191ea6f87e56559af9e02bf548509058058f6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Mislav=20Marohni=C4=87?= <mislav@github.com>
Date: Mon, 21 Jun 2021 16:22:10 +0200
Subject: [PATCH] Enable setting an alias for an extension command

---
 pkg/cmd/alias/set/set.go | 39 ++++++++++++++++++++++++---------------
 1 file changed, 24 insertions(+), 15 deletions(-)

diff --git a/pkg/cmd/alias/set/set.go b/pkg/cmd/alias/set/set.go
index e50ca3f86..48b1231be 100644
--- a/pkg/cmd/alias/set/set.go
+++ b/pkg/cmd/alias/set/set.go
@@ -20,7 +20,8 @@ type SetOptions struct {
 	Name      string
 	Expansion string
 	IsShell   bool
-	RootCmd   *cobra.Command
+
+	validCommand func(string) bool
 }
 
 func NewCmdSet(f *cmdutil.Factory, runF func(*SetOptions) error) *cobra.Command {
@@ -78,11 +79,29 @@ func NewCmdSet(f *cmdutil.Factory, runF func(*SetOptions) error) *cobra.Command
 		`),
 		Args: cobra.ExactArgs(2),
 		RunE: func(cmd *cobra.Command, args []string) error {
-			opts.RootCmd = cmd.Root()
-
 			opts.Name = args[0]
 			opts.Expansion = args[1]
 
+			opts.validCommand = func(args string) bool {
+				split, err := shlex.Split(args)
+				if err != nil {
+					return false
+				}
+
+				rootCmd := cmd.Root()
+				cmd, _, err := rootCmd.Traverse(split)
+				if err == nil && cmd != rootCmd {
+					return true
+				}
+
+				for _, ext := range f.ExtensionManager.List() {
+					if ext.Name() == split[0] {
+						return true
+					}
+				}
+				return false
+			}
+
 			if runF != nil {
 				return runF(opts)
 			}
@@ -143,11 +162,11 @@ func setRun(opts *SetOptions) error {
 	}
 	isShell = strings.HasPrefix(expansion, "!")
 
-	if validCommand(opts.RootCmd, opts.Name) {
+	if opts.validCommand(opts.Name) {
 		return fmt.Errorf("could not create alias: %q is already a gh command", opts.Name)
 	}
 
-	if !isShell && !validCommand(opts.RootCmd, expansion) {
+	if !isShell && !opts.validCommand(expansion) {
 		return fmt.Errorf("could not create alias: %s does not correspond to a gh command", expansion)
 	}
 
@@ -173,16 +192,6 @@ func setRun(opts *SetOptions) error {
 	return nil
 }
 
-func validCommand(rootCmd *cobra.Command, expansion string) bool {
-	split, err := shlex.Split(expansion)
-	if err != nil {
-		return false
-	}
-
-	cmd, _, err := rootCmd.Traverse(split)
-	return err == nil && cmd != rootCmd
-}
-
 func getExpansion(opts *SetOptions) (string, error) {
 	if opts.Expansion == "-" {
 		stdin, err := ioutil.ReadAll(opts.IO.In)