Merge pull request #100 from eparis/pre-post-run

Pre post run fucntions

README.md

@@ -348,6 +348,76 @@ Like help the function and template are over ridable through public methods.
 
     command.SetUsageTemplate(s string)
 
+## PreRun or PostRun Hooks
+
+It is possible to run functions before or after the main `Run` function of your command. The `PersistentPreRun` and `PreRun` functions will be executed before `Run`. `PersistendPostRun` and `PostRun` will be executed after `Run`.  The `Persistent*Run` functions will be inherrited by children if they do not declare their own.  These function are run in the following order:
+
+- `PersistentPreRun`
+- `PreRun`
+- `Run`
+- `PostRun`
+- `PersistenPostRun`
+
+And example of two commands which use all of these features is below.  When the subcommand in executed it will run the root command's `PersistentPreRun` but not the root command's `PersistentPostRun`
+
+```go
+package main
+
+import (
+	"fmt"
+
+	"github.com/spf13/cobra"
+)
+
+func main() {
+
+	var rootCmd = &cobra.Command{
+		Use:   "root [sub]",
+		Short: "My root command",
+		PersistentPreRun: func(cmd *cobra.Command, args []string) {
+			fmt.Printf("Inside rootCmd PersistentPreRun with args: %v\n", args)
+		},
+		PreRun: func(cmd *cobra.Command, args []string) {
+			fmt.Printf("Inside rootCmd PreRun with args: %v\n", args)
+		},
+		Run: func(cmd *cobra.Command, args []string) {
+			fmt.Printf("Inside rootCmd Run with args: %v\n", args)
+		},
+		PostRun: func(cmd *cobra.Command, args []string) {
+			fmt.Printf("Inside rootCmd PostRun with args: %v\n", args)
+		},
+		PersistentPostRun: func(cmd *cobra.Command, args []string) {
+			fmt.Printf("Inside rootCmd PersistentPostRun with args: %v\n", args)
+		},
+	}
+
+	var subCmd = &cobra.Command{
+		Use:   "sub [no options!]",
+		Short: "My sub command",
+		PreRun: func(cmd *cobra.Command, args []string) {
+			fmt.Printf("Inside subCmd PreRun with args: %v\n", args)
+		},
+		Run: func(cmd *cobra.Command, args []string) {
+			fmt.Printf("Inside subCmd Run with args: %v\n", args)
+		},
+		PostRun: func(cmd *cobra.Command, args []string) {
+			fmt.Printf("Inside subCmd PostRun with args: %v\n", args)
+		},
+		PersistentPostRun: func(cmd *cobra.Command, args []string) {
+			fmt.Printf("Inside subCmd PersistentPostRun with args: %v\n", args)
+		},
+	}
+
+	rootCmd.AddCommand(subCmd)
+
+	rootCmd.SetArgs([]string{""})
+	_ = rootCmd.Execute()
+	fmt.Print("\n")
+	rootCmd.SetArgs([]string{"sub", "arg1", "arg2"})
+	_ = rootCmd.Execute()
+}
+```
+
 ## Generating markdown formatted documentation for your command
 
 Cobra can generate a markdown formatted document based on the subcommands, flags, etc. A simple example of how to do this for your command can be found in [Markdown Docs](md_docs.md)

cobra_test.go

@@ -13,6 +13,7 @@ var _ = fmt.Println
 var _ = os.Stderr
 
 var tp, te, tt, t1 []string
+var rootPersPre, echoPre, echoPersPre, timesPersPre []string
 var flagb1, flagb2, flagb3, flagbr, flagbp bool
 var flags1, flags2a, flags2b, flags3 string
 var flagi1, flagi2, flagi3, flagir int
@@ -38,6 +39,12 @@ var cmdEcho = &Command{
 	Short:   "Echo anything to the screen",
 	Long:    `an utterly useless command for testing.`,
 	Example: "Just run cobra-test echo",
+	PersistentPreRun: func(cmd *Command, args []string) {
+		echoPersPre = args
+	},
+	PreRun: func(cmd *Command, args []string) {
+		echoPre = args
+	},
 	Run: func(cmd *Command, args []string) {
 		te = args
 	},
@@ -64,6 +71,9 @@ var cmdTimes = &Command{
 	Use:   "times [# times] [string to echo]",
 	Short: "Echo anything to the screen more times",
 	Long:  `a slightly useless command for testing.`,
+	PersistentPreRun: func(cmd *Command, args []string) {
+		timesPersPre = args
+	},
 	Run: func(cmd *Command, args []string) {
 		tt = args
 	},
@@ -73,6 +83,9 @@ var cmdRootNoRun = &Command{
 	Use:   "cobra-test",
 	Short: "The root can run it's own function",
 	Long:  "The root description for help",
+	PersistentPreRun: func(cmd *Command, args []string) {
+		rootPersPre = args
+	},
 }
 
 var cmdRootSameName = &Command{
@@ -149,6 +162,8 @@ func commandInit() {
 
 func initialize() *Command {
 	tt, tp, te = nil, nil, nil
+	rootPersPre, echoPre, echoPersPre, timesPersPre = nil, nil, nil, nil
+
 	var c = cmdRootNoRun
 	flagInit()
 	commandInit()
@@ -157,6 +172,7 @@ func initialize() *Command {
 
 func initializeWithSameName() *Command {
 	tt, tp, te = nil, nil, nil
+	rootPersPre, echoPre, echoPersPre, timesPersPre = nil, nil, nil, nil
 	var c = cmdRootSameName
 	flagInit()
 	commandInit()
@@ -827,3 +843,29 @@ func TestDeprecatedSub(t *testing.T) {
 
 	checkResultContains(t, c, cmdDeprecated.Deprecated)
 }
+
+func TestPreRun(t *testing.T) {
+	noRRSetupTest("echo one two")
+	if echoPre == nil || echoPersPre == nil {
+		t.Error("PreRun or PersistentPreRun not called")
+	}
+	if rootPersPre != nil || timesPersPre != nil {
+		t.Error("Wrong *Pre functions called!")
+	}
+
+	noRRSetupTest("echo times one two")
+	if timesPersPre == nil {
+		t.Error("PreRun or PersistentPreRun not called")
+	}
+	if echoPre != nil || echoPersPre != nil || rootPersPre != nil {
+		t.Error("Wrong *Pre functions called!")
+	}
+
+	noRRSetupTest("print one two")
+	if rootPersPre == nil {
+		t.Error("Parent PersistentPreRun not called but should not have been")
+	}
+	if echoPre != nil || echoPersPre != nil || timesPersPre != nil {
+		t.Error("Wrong *Pre functions called!")
+	}
+}

command.go

@@ -56,9 +56,23 @@ type Command struct {
 	pflags *flag.FlagSet
 	// Flags that are declared specifically by this command (not inherited).
 	lflags *flag.FlagSet
-	// Run runs the command.
+	// The *Run functions are executed in the following order:
-	// The args are the arguments after the command name.
+	//   * PersistentPreRun()
+	//   * PreRun()
+	//   * Run()
+	//   * PostRun()
+	//   * PersistentPostRun()
+	// All functions get the same args, the arguments after the command name
+	// PersistentPreRun: children of this command will inherit and execute
+	PersistentPreRun func(cmd *Command, args []string)
+	// PreRun: children of this command will not inherit.
+	PreRun func(cmd *Command, args []string)
+	// Run: Typically the actual work function. Most commands will only implement this
 	Run func(cmd *Command, args []string)
+	// PostRun: run after the Run command.
+	PostRun func(cmd *Command, args []string)
+	// PersistentPostRun: children of this command will inherit and execute after PostRun
+	PersistentPostRun func(cmd *Command, args []string)
 	// Commands is the list of commands supported by this program.
 	commands []*Command
 	// Parent Command for this command
@@ -448,7 +462,23 @@ func (c *Command) execute(a []string) (err error) {
 
 	c.preRun()
 	argWoFlags := c.Flags().Args()
+
+	if c.PersistentPreRun != nil {
+		c.PersistentPreRun(c, argWoFlags)
+	}
+	if c.PreRun != nil {
+		c.PreRun(c, argWoFlags)
+	}
+
 	c.Run(c, argWoFlags)
+
+	if c.PostRun != nil {
+		c.PostRun(c, argWoFlags)
+	}
+	if c.PersistentPostRun != nil {
+		c.PersistentPostRun(c, argWoFlags)
+	}
+
 	return nil
 }
 
@@ -530,6 +560,8 @@ func (c *Command) initHelp() {
 			Long: `Help provides help for any command in the application.
     Simply type ` + c.Name() + ` help [path to command] for full details.`,
 			Run:               c.HelpFunc(),
+			PersistentPreRun:  func(cmd *Command, args []string) {},
+			PersistentPostRun: func(cmd *Command, args []string) {},
 		}
 	}
 	c.AddCommand(c.helpCommand)
@@ -569,6 +601,14 @@ func (c *Command) AddCommand(cmds ...*Command) {
 			c.commandsMaxNameLen = nameLen
 		}
 		c.commands = append(c.commands, x)
+
+		// Pass on peristent pre/post functions to children
+		if x.PersistentPreRun == nil {
+			x.PersistentPreRun = c.PersistentPreRun
+		}
+		if x.PersistentPostRun == nil {
+			x.PersistentPostRun = c.PersistentPostRun
+		}
 	}
 }