Specify number of splits to make for command arguments

bot/command.go

@@ -21,6 +21,7 @@ type (
 		Name    string
 		Config  Config
 		Func    func(cmd *Command, args []string) error
+		NArgs   int
 		Session *discordgo.Session
 		Message *discordgo.MessageCreate
 	}
@@ -66,22 +67,25 @@ func NewCommandHandler(config Config) func(s *discordgo.Session, m *discordgo.Me
 			return
 		}
 
-		cmdName, args := lib.SplitCommandAndArgs(m.Content, config.Prefix)
+		cmdName, arg := lib.SplitCommandAndArg(m.Content, config.Prefix)
 
 		cmd, ok := GetCommand(cmdName)
+
+		args := lib.SplitArgs(arg, cmd.NArgs)
+
 		if ok {
 			cmd.Config = config
 			cmd.Name = cmdName
 			cmd.Session = s
 			cmd.Message = m
 
-			log.Debugf("command: %+v, args: %+v", cmd.Name, args)
+			log.Debugf("command: %v, args: %v, nargs: %d", cmd.Name, args, len(args))
 			cmd.Func(cmd, args)
 
 			return
 		}
 
-		log.Warnf("unknown command: %+v, args: %+v", cmdName, args)
+		log.Warnf("unknown command: %v, args: %v, nargs: %d", cmdName, args, len(args))
 		s.ChannelMessageSend(m.ChannelID, fmt.Sprintf("unknown command: %s", cmdName))
 	}
 }

cmd/bb/main.go

@@ -37,6 +37,7 @@ func init() {
 	bot.AddCommand(&bot.Command{
 		Name:  "deal",
 		Func:  commands.DealCommand,
+		NArgs: 1,
 	})
 	bot.AddCommand(&bot.Command{
 		Name: "ping",
@@ -45,10 +46,12 @@ func init() {
 	bot.AddCommand(&bot.Command{
 		Name:  "roll",
 		Func:  commands.RollCommand,
+		NArgs: 1,
 	})
 	bot.AddCommand(&bot.Command{
 		Name:  "time",
 		Func:  commands.TimeCommand,
+		NArgs: 1,
 	})
 	bot.AddCommand(&bot.Command{
 		Name: "version",
@@ -57,6 +60,7 @@ func init() {
 	bot.AddCommand(&bot.Command{
 		Name:  "weather",
 		Func:  commands.WeatherCommand,
+		NArgs: 1,
 	})
 }
 

lib/common.go

@@ -99,24 +99,43 @@ func ContainsCommand(s, prefix, cmd string) bool {
 	return false
 }
 
-func SplitCommandAndArg(s, prefix string) (cmd string, args []string) {
+func SplitCommandAndArg(s, prefix string) (cmd string, arg string) {
 	s = strings.TrimSpace(s)
 
 	if !strings.HasPrefix(s, prefix) {
 		return
 	}
 
-	x := strings.Split(s, " ")
+	// remove the command prefix
+	s = s[len(prefix):]
 
+	// multiple assignment trick
+	cmd, arg = func() (string, string) {
+		x := strings.SplitN(s, " ", 2)
 		if len(x) > 1 {
-		args = x[1:]
+			return x[0], x[1]
+		}
+		return x[0], ""
+	}()
+
+	return cmd, arg
 }
 
-	cmd = x[0]
+func SplitCommandAndArgs(s, prefix string, n int) (cmd string, args []string) {
+	cmd, arg := SplitCommandAndArg(s, prefix)
 
-	if strings.Index(s, prefix) == 0 {
-		cmd = cmd[len(prefix):]
+	if n == 0 {
+		return cmd, strings.Split(arg, " ")
 	}
 
-	return cmd, args
+	return cmd, strings.SplitN(arg, " ", n)
+}
+
+func SplitArgs(s string, n int) (args []string) {
+	if n > 0 {
+		args = strings.SplitN(s, " ", n)
+	} else {
+		args = strings.Split(s, " ")
+	}
+	return
 }

lib/common_test.go

@@ -57,20 +57,44 @@ func TestHasCommandCommand(t *testing.T) {
 	}
 }
 
-func TestSplitComandAndArgs(t *testing.T) {
+func TestSplitCommandAndArg(t *testing.T) {
 	tables := []struct {
 		s       string
 		prefix  string
 		wantCmd string
-		wantArgs []string
+		wantArg string
 	}{
-		{"!command x y", "!", "command", []string{"x", "y"}},
-		{"!command", "!", "command", []string(nil)},
-		{"hey man", "!", "", []string(nil)},
+		{"!command x y", "!", "command", "x y"},
+		{"!command", "!", "command", ""},
+		{"hey man", "!", "", ""},
 	}
 
 	for _, table := range tables {
-		gotCmd, gotArgs := SplitCommandAndArgs(table.s, table.prefix)
+		gotCmd, gotArg := SplitCommandAndArg(table.s, table.prefix)
+		if gotCmd != table.wantCmd {
+			t.Errorf("got: %s, want: %s", gotCmd, table.wantCmd)
+		}
+		if gotArg != table.wantArg {
+			t.Errorf("got: %+v, want: %+v", gotArg, table.wantArg)
+		}
+	}
+}
+
+func TestSplitCommandAndArgs(t *testing.T) {
+	tables := []struct {
+		s        string
+		prefix   string
+		n        int
+		wantCmd  string
+		wantArgs []string
+	}{
+		{"!command x y", "!", 2, "command", []string{"x", "y"}},
+		{"!command x y z", "!", 2, "command", []string{"x", "y z"}},
+		{"!command", "!", 1, "command", []string{""}},
+		{"hey man", "!", 1, "", []string{""}},
+	}
+	for _, table := range tables {
+		gotCmd, gotArgs := SplitCommandAndArgs(table.s, table.prefix, table.n)
 		if gotCmd != table.wantCmd {
 			t.Errorf("got: %s, want: %s", gotCmd, table.wantCmd)
 		}
@@ -79,3 +103,22 @@ func TestSplitComandAndArgs(t *testing.T) {
 		}
 	}
 }
+
+func TestSplitArgs(t *testing.T) {
+	tables := []struct {
+		s    string
+		n    int
+		want []string
+	}{
+		{"a b c", 0, []string{"a", "b", "c"}},
+		{"a b c", 1, []string{"a b c"}},
+		{"a b c", 2, []string{"a", "b c"}},
+		{"a b c", 3, []string{"a", "b", "c"}},
+		{"a b c", 4, []string{"a", "b", "c"}},
+	}
+	for _, table := range tables {
+		if got, want := SplitArgs(table.s, table.n), table.want; !reflect.DeepEqual(got, want) {
+			t.Errorf("got: %#v, want: %#v", got, want)
+		}
+	}
+}