Add arg limiting special comment

cv · cv · commit fb8d732d7a5f · 2018-04-13T17:06:31.000-03:00
diff --git a/README.md b/README.md
@@ -81,6 +81,7 @@ In order to document the script, `sd` pays attention to a few special comments:
 # bar: Bars the foos.
 #
 # example: foo bar 123
+# args: 3
 #
 echo "sd foo bar has been called"
 ```
diff --git a/cli/cli.go b/cli/cli.go
@@ -7,6 +7,7 @@ import (
 	"os"
 	"path/filepath"
 	"regexp"
+	"strconv"
 	"strings"
 	"syscall"
 
@@ -259,6 +260,43 @@ func shortDescriptionFrom(path string) (string, error) {
 	return "", nil
 }
 
+/*
+Argument validators. Looks for a line like this:
+
+# args: 3
+
+ */
+func argsFrom(path string) (cobra.PositionalArgs, error) {
+	file, err := os.Open(path)
+	if err != nil {
+		return nil, err
+	}
+	defer func() {
+		err = file.Close()
+		if err != nil {
+			logrus.Error(err)
+		}
+	}()
+
+	r := regexp.MustCompile(`^# args: (\d+)$`)
+	scanner := bufio.NewScanner(file)
+	for scanner.Scan() {
+		match := r.FindStringSubmatch(scanner.Text())
+		if len(match) == 2 {
+			logrus.Debug("Found example line: ", filepath.Join(path), ", set to: ", match[1])
+			i, err := strconv.ParseUint(match[1], 10, 32)
+			if err != nil {
+				return nil, err
+			}
+			if i == 0 {
+				return cobra.NoArgs, nil
+			}
+			return cobra.ExactArgs(int(i)), nil
+		}
+	}
+	return cobra.ArbitraryArgs, nil
+}
+
 /*
 
 Looks for a line like this:
@@ -301,13 +339,19 @@ func commandFromScript(path string) (*cobra.Command, error) {
 		return nil, err
 	}
 
+	args, err := argsFrom(path)
+	if err != nil {
+		return nil, err
+	}
+
 	cmd := &cobra.Command{
 		Use:     filepath.Base(path),
 		Short:   shortDesc,
 		Example: example,
 		Annotations: map[string]string{
 			"Source": path,
 		},
+		Args: args,
 		RunE: execCommand,
 	}
 
@@ -318,7 +362,7 @@ func commandFromScript(path string) (*cobra.Command, error) {
 // these get mocked in tests
 var (
 	syscallExec = syscall.Exec
-	env = os.Getenv
+	env         = os.Getenv
 )
 
 func execCommand(cmd *cobra.Command, args []string) error {
diff --git a/cli/cli_test.go b/cli/cli_test.go
@@ -152,6 +152,92 @@ func TestShortDescriptionFrom(t *testing.T) {
 	})
 }
 
+func TestArgsFrom(t *testing.T) {
+	t.Run("happy path with 3 args", func(t *testing.T) {
+		f, err := ioutil.TempFile("", "test-args3-ok")
+		assert.NoError(t, err)
+
+		f.WriteString("#\n# args: 2\n#\n")
+		defer func() {
+			f.Close()
+			os.Remove(f.Name())
+		}()
+
+		v, err := argsFrom(f.Name())
+		assert.NoError(t, err)
+		assert.Error(t, v(&cobra.Command{}, []string{}))
+		assert.Error(t, v(&cobra.Command{}, []string{"first"}))
+		assert.NoError(t, v(&cobra.Command{}, []string{"first", "second"}))
+		assert.Error(t, v(&cobra.Command{}, []string{"first", "second", "third"}))
+	})
+
+	t.Run("happy path with no args", func(t *testing.T) {
+		f, err := ioutil.TempFile("", "test-args0-ok")
+		assert.NoError(t, err)
+
+		f.WriteString("#\n# args: 0\n#\n")
+		defer func() {
+			f.Close()
+			os.Remove(f.Name())
+		}()
+
+		v, err := argsFrom(f.Name())
+		assert.NoError(t, err)
+		assert.NoError(t, v(&cobra.Command{}, []string{}))
+		assert.Error(t, v(&cobra.Command{}, []string{"first"}))
+		assert.Error(t, v(&cobra.Command{}, []string{"first", "second"}))
+	})
+
+	t.Run("missing", func(t *testing.T) {
+		f, err := ioutil.TempFile("", "test-args-missing")
+		assert.NoError(t, err)
+
+		f.WriteString("#\n#\n#\n")
+		defer func() {
+			f.Close()
+			os.Remove(f.Name())
+		}()
+
+		v, err := argsFrom(f.Name())
+		assert.NoError(t, err)
+		assert.NoError(t, v(&cobra.Command{}, []string{}))
+		assert.NoError(t, v(&cobra.Command{}, []string{"first", "second"}))
+	})
+
+	t.Run("bad value: not numeric", func(t *testing.T) {
+		f, err := ioutil.TempFile("", "test-args-not-numeric")
+		assert.NoError(t, err)
+
+		f.WriteString("#\n# args: banana\n#\n")
+		defer func() {
+			f.Close()
+			os.Remove(f.Name())
+		}()
+
+		v, err := argsFrom(f.Name())
+		assert.NoError(t, err)
+		assert.NoError(t, v(&cobra.Command{}, []string{}))
+		assert.NoError(t, v(&cobra.Command{}, []string{"first", "second"}))
+	})
+
+	t.Run("bad value: too large", func(t *testing.T) {
+		f, err := ioutil.TempFile("", "test-args-too-large")
+		assert.NoError(t, err)
+
+		f.WriteString("#\n# args: 999999999999999999999999999999999999999999999999999999999999999999999999\n#\n")
+		defer func() {
+			f.Close()
+			os.Remove(f.Name())
+		}()
+
+		v, err := argsFrom(f.Name())
+		if !assert.Error(t, err) {
+			assert.NoError(t, v(&cobra.Command{}, []string{}))
+			assert.NoError(t, v(&cobra.Command{}, []string{"first", "second"}))
+		}
+	})
+}
+
 func TestExampleFrom(t *testing.T) {
 	t.Run("happy path", func(t *testing.T) {
 		f, err := ioutil.TempFile("", "test-example-ok")

Original file line number	Diff line number	Diff line change
@@ -81,6 +81,7 @@ In order to document the script, `sd` pays attention to a few special comments:
`81`	`81`	`# bar: Bars the foos.`
`82`	`82`	`#`
`83`	`83`	`# example: foo bar 123`
	`84`	`+# args: 3`
`84`	`85`	`#`
`85`	`86`	`echo "sd foo bar has been called"`
`86`	`87`	```