From 139b32094e93324971bd9d720431716ae6a9b526 Mon Sep 17 00:00:00 2001 From: Ryan Cavicchioni Date: Tue, 6 Sep 2022 10:43:30 -0500 Subject: [PATCH] Specify number of splits to make for command arguments --- bot/command.go | 10 +++++--- cmd/bb/main.go | 20 +++++++++------ lib/common.go | 43 +++++++++++++++++++++++--------- lib/common_test.go | 61 +++++++++++++++++++++++++++++++++++++++------- 4 files changed, 102 insertions(+), 32 deletions(-) diff --git a/bot/command.go b/bot/command.go index 92f796f..be3db97 100644 --- a/bot/command.go +++ b/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)) } } diff --git a/cmd/bb/main.go b/cmd/bb/main.go index d645bf3..78b4f27 100644 --- a/cmd/bb/main.go +++ b/cmd/bb/main.go @@ -35,28 +35,32 @@ func init() { Func: commands.CoinCommand, }) bot.AddCommand(&bot.Command{ - Name: "deal", - Func: commands.DealCommand, + Name: "deal", + Func: commands.DealCommand, + NArgs: 1, }) bot.AddCommand(&bot.Command{ Name: "ping", Func: commands.PingCommand, }) bot.AddCommand(&bot.Command{ - Name: "roll", - Func: commands.RollCommand, + Name: "roll", + Func: commands.RollCommand, + NArgs: 1, }) bot.AddCommand(&bot.Command{ - Name: "time", - Func: commands.TimeCommand, + Name: "time", + Func: commands.TimeCommand, + NArgs: 1, }) bot.AddCommand(&bot.Command{ Name: "version", Func: commands.VersionCommand, }) bot.AddCommand(&bot.Command{ - Name: "weather", - Func: commands.WeatherCommand, + Name: "weather", + Func: commands.WeatherCommand, + NArgs: 1, }) } diff --git a/lib/common.go b/lib/common.go index 8f69cd3..d7de19e 100644 --- a/lib/common.go +++ b/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):] - if len(x) > 1 { - args = x[1:] - } + // multiple assignment trick + cmd, arg = func() (string, string) { + x := strings.SplitN(s, " ", 2) + if len(x) > 1 { + return x[0], x[1] + } + return x[0], "" + }() - cmd = x[0] - - if strings.Index(s, prefix) == 0 { - cmd = cmd[len(prefix):] - } - - return cmd, args + return cmd, arg +} + +func SplitCommandAndArgs(s, prefix string, n int) (cmd string, args []string) { + cmd, arg := SplitCommandAndArg(s, prefix) + + if n == 0 { + return cmd, strings.Split(arg, " ") + } + + 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 } diff --git a/lib/common_test.go b/lib/common_test.go index ff836d9..b262655 100644 --- a/lib/common_test.go +++ b/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 + s string + prefix string + wantCmd 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) + } + } +}