Implement command groups and buckets

* Implement command groups * change to ref mut * Implement framework API. * Remove commands field * Make it all work * Make example use command groups * Requested changes * Implement adding buckets * Add ratelimit check function * Finish everything * Fix voice example * Actually fix it * Fix doc tests * Switch to result * Savage examples * Fix docs * Fixes * Accidental push * 👀 * Fix an example * fix some example * Small cleanup * Abstract ratelimit bucket logic
serenity-rs · Dec 13, 2016 · daf92ed · daf92ed
1 parent fbe5600
commit daf92ed
Show file tree

Hide file tree

Showing 14 changed files with 775 additions and 131 deletions.
diff --git a/README.md b/README.md
@@ -55,9 +55,9 @@ fn main() {
     let _ = client.start();
 }
 
-fn ping(_context: &Context, message: &Message, _args: Vec<String>) {
+command!(ping(_context, message) {
     let _ = message.reply("Pong!");
-}
+});
 ```
 
 ### Full Examples

diff --git a/examples/06_command_framework/src/main.rs b/examples/06_command_framework/src/main.rs
@@ -15,6 +15,7 @@ extern crate typemap;
 use serenity::client::Context;
 use serenity::Client;
 use serenity::model::{Message, permissions};
+use serenity::ext::framework::help_commands;
 use std::collections::HashMap;
 use std::env;
 use std::fmt::Write;
@@ -58,6 +59,7 @@ fn main() {
         .configure(|c| c
             .allow_whitespace(true)
             .on_mention(true)
+            .rate_limit_message("Try this again in `%time%` seconds.")
             .prefix("~"))
         // Set a function to be called prior to each command execution. This
         // provides the context of the command, the message that was received,
@@ -81,17 +83,35 @@ fn main() {
         })
         // Very similar to `before`, except this will be called directly _after_
         // command execution.
-        .after(|_, _, command_name| {
-            println!("Processed command '{}'", command_name)
+        .after(|_, _, command_name, error| {
+            if let Some(why) = error {
+                println!("Command '{}' returned error {:?}", command_name, why);
+            } else {
+                println!("Processed command '{}'", command_name);
+            }
         })
+        // Can't be used more than once per 5 seconds:
+        .simple_bucket("emoji", 5)
+        // Can't be used more than 2 times per 30 seconds, with a 5 second delay
+        .bucket("complicated", 5, 30, 2)
         .command("about", |c| c.exec_str("A test bot"))
+        .command("help", |c| c.exec_help(help_commands::plain))
         .command("commands", |c| c
-            .check(owner_check)
+            // Make this command use the "complicated" bucket.
+            .bucket("complicated")
             .exec(commands))
-        .command("emoji cat", |c| c
-            .exec_str(":cat:")
-            .required_permissions(permissions::SEND_MESSAGES))
-        .command("emoji dog", |c| c.exec_str(":dog:"))
+        .group("Emoji", |g| g
+            .prefix("emoji")
+            .command("cat", |c| c
+                .desc("Sends an emoji with a cat.")
+                .bucket("emoji") // Make this command use the "emoji" bucket.
+                .exec_str(":cat:")
+                 // Allow only administrators to call this:
+                .required_permissions(permissions::ADMINISTRATOR))
+            .command("dog", |c| c
+                .desc("Sends an emoji with a dog.")
+                .bucket("emoji")
+                .exec_str(":dog:")))
         .command("multiply", |c| c.exec(multiply))
         .command("ping", |c| c
             .check(owner_check)
@@ -134,11 +154,11 @@ fn owner_check(_: &Context, message: &Message) -> bool {
     message.author.id == 7
 }
 
-fn some_long_command(context: &Context, _: &Message, args: Vec<String>) {
+command!(some_long_command(context, _msg, args) {
     if let Err(why) = context.say(&format!("Arguments: {:?}", args)) {
         println!("Error sending message: {:?}", why);
     }
-}
+});
 
 // Using the `command!` macro, commands can be created with a certain type of
 // "dynamic" type checking. This is a method of requiring that the arguments

diff --git a/examples/07_voice/src/main.rs b/examples/07_voice/src/main.rs
@@ -7,11 +7,12 @@
 //! features = ["cache", "framework", "methods", "voice"]
 //! ```
 
+#[macro_use]
 extern crate serenity;
 
 use serenity::client::{CACHE, Client, Context};
 use serenity::ext::voice;
-use serenity::model::{ChannelId, Message};
+use serenity::model::{ChannelId, Message, Mentionable};
 use serenity::Result as SerenityResult;
 use std::env;
 
@@ -41,13 +42,13 @@ fn main() {
     let _ = client.start().map_err(|why| println!("Client ended: {:?}", why));
 }
 
-fn deafen(context: &Context, message: &Message, _args: Vec<String>) {
+command!(deafen(context, message) {
     let guild_id = match CACHE.read().unwrap().get_guild_channel(message.channel_id) {
         Some(channel) => channel.guild_id,
         None => {
             check_msg(context.say("Groups and DMs not supported"));
 
-            return;
+            return Ok(());
         },
     };
 
@@ -58,7 +59,7 @@ fn deafen(context: &Context, message: &Message, _args: Vec<String>) {
         None => {
             check_msg(message.reply("Not in a voice channel"));
 
-            return;
+            return Ok(());
         },
     };
 
@@ -69,22 +70,22 @@ fn deafen(context: &Context, message: &Message, _args: Vec<String>) {
 
         check_msg(context.say("Deafened"));
     }
-}
+});
 
-fn join(context: &Context, message: &Message, args: Vec<String>) {
+command!(join(context, message, args) {
     let connect_to = match args.get(0) {
         Some(arg) => match arg.parse::<u64>() {
             Ok(id) => ChannelId(id),
             Err(_why) => {
                 check_msg(message.reply("Invalid voice channel ID given"));
 
-                return;
+                return Ok(());
             },
         },
         None => {
             check_msg(message.reply("Requires a voice channel ID be given"));
 
-            return;
+            return Ok(());
         },
     };
 
@@ -93,23 +94,23 @@ fn join(context: &Context, message: &Message, args: Vec<String>) {
         None => {
             check_msg(context.say("Groups and DMs not supported"));
 
-            return;
+            return Ok(());
         },
     };
 
     let mut shard = context.shard.lock().unwrap();
     shard.manager.join(Some(guild_id), connect_to);
 
     check_msg(context.say(&format!("Joined {}", connect_to.mention())));
-}
+});
 
-fn leave(context: &Context, message: &Message, _args: Vec<String>) {
+command!(leave(context, message) {
     let guild_id = match CACHE.read().unwrap().get_guild_channel(message.channel_id) {
         Some(channel) => channel.guild_id,
         None => {
             check_msg(context.say("Groups and DMs not supported"));
 
-            return;
+            return Ok(());
         },
     };
 
@@ -123,15 +124,15 @@ fn leave(context: &Context, message: &Message, _args: Vec<String>) {
     } else {
         check_msg(message.reply("Not in a voice channel"));
     }
-}
+});
 
-fn mute(context: &Context, message: &Message, _args: Vec<String>) {
+command!(mute(context, message) {
     let guild_id = match CACHE.read().unwrap().get_guild_channel(message.channel_id) {
         Some(channel) => channel.guild_id,
         None => {
             check_msg(context.say("Groups and DMs not supported"));
 
-            return;
+            return Ok(());
         },
     };
 
@@ -142,7 +143,7 @@ fn mute(context: &Context, message: &Message, _args: Vec<String>) {
         None => {
             check_msg(message.reply("Not in a voice channel"));
 
-            return;
+            return Ok(());
         },
     };
 
@@ -153,34 +154,34 @@ fn mute(context: &Context, message: &Message, _args: Vec<String>) {
 
         check_msg(context.say("Now muted"));
     }
-}
+});
 
-fn ping(context: &Context, _message: &Message, _args: Vec<String>) {
+command!(ping(context) {
     check_msg(context.say("Pong!"));
-}
+});
 
-fn play(context: &Context, message: &Message, args: Vec<String>) {
+command!(play(context, message, args) {
     let url = match args.get(0) {
         Some(url) => url,
         None => {
             check_msg(context.say("Must provide a URL to a video or audio"));
 
-            return;
+            return Ok(());
         },
     };
 
     if !url.starts_with("http") {
         check_msg(context.say("Must provide a valid URL"));
 
-        return;
+        return Ok(());
     }
 
     let guild_id = match CACHE.read().unwrap().get_guild_channel(message.channel_id) {
         Some(channel) => channel.guild_id,
         None => {
             check_msg(context.say("Error finding channel info"));
 
-            return;
+            return Ok(());
         },
     };
 
@@ -192,7 +193,7 @@ fn play(context: &Context, message: &Message, args: Vec<String>) {
 
                 check_msg(context.say("Error sourcing ffmpeg"));
 
-                return;
+                return Ok(());
             },
         };
 
@@ -202,15 +203,15 @@ fn play(context: &Context, message: &Message, args: Vec<String>) {
     } else {
         check_msg(context.say("Not in a voice channel to play in"));
     }
-}
+});
 
-fn undeafen(context: &Context, message: &Message, _args: Vec<String>) {
+command!(undeafen(context, message) {
     let guild_id = match CACHE.read().unwrap().get_guild_channel(message.channel_id) {
         Some(channel) => channel.guild_id,
         None => {
             check_msg(context.say("Error finding channel info"));
 
-            return;
+            return Ok(());
         },
     };
 
@@ -221,15 +222,15 @@ fn undeafen(context: &Context, message: &Message, _args: Vec<String>) {
     } else {
         check_msg(context.say("Not in a voice channel to undeafen in"));
     }
-}
+});
 
-fn unmute(context: &Context, message: &Message, _args: Vec<String>) {
+command!(unmute(context, message) {
     let guild_id = match CACHE.read().unwrap().get_guild_channel(message.channel_id) {
         Some(channel) => channel.guild_id,
         None => {
             check_msg(context.say("Error finding channel info"));
 
-            return;
+            return Ok(());
         },
     };
 
@@ -240,10 +241,9 @@ fn unmute(context: &Context, message: &Message, _args: Vec<String>) {
     } else {
         check_msg(context.say("Not in a voice channel to undeafen in"));
     }
-}
+});
 
 /// Checks that a message successfully sent; if not, then logs why to stdout.
-#[cfg(feature="voice")]
 fn check_msg(result: SerenityResult<Message>) {
     if let Err(why) = result {
         println!("Error sending message: {:?}", why);

diff --git a/src/client/context.rs b/src/client/context.rs
@@ -902,7 +902,7 @@ impl Context {
     /// Change the current user's username:
     ///
     /// ```rust,ignore
-    /// context.edit_member(|p| p.username("meew zero"));
+    /// context.edit_profile(|p| p.username("meew0"));
     /// ```
     pub fn edit_profile<F: FnOnce(EditProfile) -> EditProfile>(&self, f: F)
         -> Result<CurrentUser> {
@@ -1534,7 +1534,7 @@ impl Context {
     ///
     /// let _ = client.start();
     ///
-    /// fn ping(context: &Context, message: &Message, _arguments: Vec<String>) {
+    /// command!(ping(context, message) {
     ///     let cache = CACHE.read().unwrap();
     ///     let channel = cache.get_guild_channel(message.channel_id);
     ///
@@ -1571,7 +1571,9 @@ impl Context {
     ///
     ///                 f
     ///             })));
-    /// }
+    ///
+    ///     Ok(())
+    /// });
     /// ```
     ///
     /// Note that for most use cases, your embed layout will _not_ be this ugly.

diff --git a/src/ext/framework/buckets.rs b/src/ext/framework/buckets.rs
@@ -0,0 +1,60 @@
+use std::collections::HashMap;
+use std::default::Default;
+use time;
+
+#[doc(hidden)]
+pub struct Ratelimit {
+    pub delay: i64,
+    pub limit: Option<(i64, i32)>,
+}
+
+#[doc(hidden)]
+pub struct MemberRatelimit {
+    pub count: i32,
+    pub last_time: i64,
+    pub set_time: i64,
+}
+
+impl Default for MemberRatelimit {
+    fn default() -> Self {
+        MemberRatelimit {
+            count: 0,
+            last_time: 0,
+            set_time: 0,
+        }
+    }
+}
+
+#[doc(hidden)]
+pub struct Bucket {
+    pub ratelimit: Ratelimit,
+    pub limits: HashMap<u64, MemberRatelimit>,
+}
+
+impl Bucket {
+    pub fn take(&mut self, user_id: u64) -> i64 {
+        let time =- time::get_time().sec;
+        let member = self.limits.entry(user_id)
+            .or_insert_with(MemberRatelimit::default);
+
+        if let Some((timespan, limit)) = self.ratelimit.limit {
+            if (member.count + 1) > limit {
+                if time < (member.set_time + timespan) {
+                    return (member.set_time + timespan) - time;
+                } else {
+                    member.count = 0;
+                    member.set_time = time;
+                }
+            }
+        }
+
+        if time < member.last_time + self.ratelimit.delay {
+            (member.last_time + self.ratelimit.delay) - time
+        } else {
+            member.count += 1;
+            member.last_time = time;
+
+            0
+        }
+    }
+}