nki/youmubot
1
0
Fork 0
mirror of https://github.com/natsukagami/youmubot.git synced 2025-04-19 08:48:54 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Merge pull request #5 from natsukagami/async-youmu

Browse source 
Async youmu!!!
This commit is contained in:
Natsu Kagami 2020-09-20 20:57:58 +00:00 • committed by GitHub
commit e25701a99c
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
41 changed files with 3471 additions and 3197 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
.gitignore vendored
View file

@ -2,3 +2,4 @@ target
.env
*.yaml
cargo-remote
.vscode

64
.vscode/launch.json vendored
View file

@ -1,64 +0,0 @@
{
// Use IntelliSense to learn about possible attributes.
// Hover to view descriptions of existing attributes.
// For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
"version": "0.2.0",
"configurations": [
{
"type": "lldb",
"request": "launch",
"name": "Debug unit tests in library 'youmubot-osu'",
"cargo": {
"args": [
"test",
"--no-run",
"--lib",
"--package=youmubot-osu"
],
"filter": {
"name": "youmubot-osu",
"kind": "lib"
}
},
"args": [],
"cwd": "${workspaceFolder}"
},
{
"type": "lldb",
"request": "launch",
"name": "Debug executable 'youmubot'",
"cargo": {
"args": [
"build",
"--bin=youmubot",
"--package=youmubot"
],
"filter": {
"name": "youmubot",
"kind": "bin"
}
},
"args": [],
"cwd": "${workspaceFolder}"
},
{
"type": "lldb",
"request": "launch",
"name": "Debug unit tests in executable 'youmubot'",
"cargo": {
"args": [
"test",
"--no-run",
"--bin=youmubot",
"--package=youmubot"
],
"filter": {
"name": "youmubot",
"kind": "bin"
}
},
"args": [],
"cwd": "${workspaceFolder}"
}
]
}

1997
Cargo.lock generated
View file

File diff suppressed because it is too large Load diff

8
youmubot-cf/Cargo.toml
View file

@ -7,15 +7,15 @@ edition = "2018"
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
[dependencies]
serde = { version = "1", features = ["derive"] }
tokio = { version = "0.2", features = ["time"] }
reqwest = "0.10.1"
serenity = "0.8"
serenity = "0.9.0-rc.0"
Inflector = "0.11"
codeforces = { git = "https://github.com/natsukagami/rust-codeforces-api" }
codeforces = "0.2.1"
regex = "1"
lazy_static = "1"
rayon = "1"
chrono = { version = "0.4", features = ["serde"] }
crossbeam-channel = "0.4"
dashmap = "3.11.4"
youmubot-prelude = { path = "../youmubot-prelude" }
youmubot-db = { path = "../youmubot-db" }

139
youmubot-cf/src/announcer.rs
View file

@ -1,85 +1,60 @@
use crate::db::{CfSavedUsers, CfUser};
use crate::{
db::{CfSavedUsers, CfUser},
CFClient,
};
use announcer::MemberToChannels;
use chrono::Utc;
use codeforces::{RatingChange, User};
use serenity::{
framework::standard::{CommandError, CommandResult},
http::CacheHttp,
model::id::{ChannelId, UserId},
CacheAndHttp,
};
use serenity::{http::CacheHttp, model::id::UserId, CacheAndHttp};
use std::sync::Arc;
use youmubot_prelude::*;
type Reqwest = <HTTPClient as TypeMapKey>::Value;
/// Updates the rating and rating changes of the users.
pub fn updates(
http: Arc<CacheAndHttp>,
data: AppData,
channels: MemberToChannels,
) -> CommandResult {
let mut users = CfSavedUsers::open(&*data.read()).borrow()?.clone();
let reqwest = data.get_cloned::<HTTPClient>();
pub struct Announcer;
for (user_id, cfu) in users.iter_mut() {
if let Err(e) = update_user(http.clone(), &channels, &reqwest, *user_id, cfu) {
dbg!((*user_id, e));
}
#[async_trait]
impl youmubot_prelude::Announcer for Announcer {
async fn updates(
&mut self,
http: Arc<CacheAndHttp>,
data: AppData,
channels: MemberToChannels,
) -> Result<()> {
let data = data.read().await;
let client = data.get::<CFClient>().unwrap();
let mut users = CfSavedUsers::open(&*data).borrow()?.clone();
users
.iter_mut()
.map(|(user_id, cfu)| update_user(http.clone(), &channels, &client, *user_id, cfu))
.collect::<stream::FuturesUnordered<_>>()
.try_collect::<()>()
.await?;
*CfSavedUsers::open(&*data).borrow_mut()? = users;
Ok(())
}
*CfSavedUsers::open(&*data.read()).borrow_mut()? = users;
Ok(())
}
fn update_user(
async fn update_user(
http: Arc<CacheAndHttp>,
channels: &MemberToChannels,
reqwest: &Reqwest,
client: &codeforces::Client,
user_id: UserId,
cfu: &mut CfUser,
) -> CommandResult {
// Ensure this takes 200ms
let after = crossbeam_channel::after(std::time::Duration::from_secs_f32(0.2));
let info = User::info(reqwest, &[cfu.handle.as_str()])?
) -> Result<()> {
let info = User::info(client, &[cfu.handle.as_str()])
.await?
.into_iter()
.next()
.ok_or(CommandError::from("Not found"))?;
.ok_or(Error::msg("Not found"))?;
let rating_changes = info.rating_changes(reqwest)?;
let rating_changes = info.rating_changes(client).await?;
let mut channels_list: Option<Vec<ChannelId>> = None;
let channels_list = channels.channels_of(&http, user_id).await;
cfu.last_update = Utc::now();
// Update the rating
cfu.rating = info.rating;
let mut send_message = |rc: RatingChange| -> CommandResult {
let channels =
channels_list.get_or_insert_with(|| channels.channels_of(http.clone(), user_id));
if channels.is_empty() {
return Ok(());
}
let (contest, _, _) =
codeforces::Contest::standings(reqwest, rc.contest_id, |f| f.limit(1, 1))?;
for channel in channels {
if let Err(e) = channel.send_message(http.http(), |e| {
e.content(format!("Rating change for {}!", user_id.mention()))
.embed(|c| {
crate::embed::rating_change_embed(
&rc,
&info,
&contest,
&user_id.mention(),
c,
)
})
}) {
dbg!(e);
}
}
Ok(())
};
let rating_changes = match cfu.last_contest_id {
None => rating_changes,
Some(v) => {
@ -101,12 +76,46 @@ fn update_user(
.or(cfu.last_contest_id);
// Check for any good announcements to make
for rc in rating_changes {
if let Err(v) = send_message(rc) {
dbg!(v);
}
}
after.recv().ok();
rating_changes
.into_iter()
.map(|rc: RatingChange| {
let channels = channels_list.clone();
let http = http.clone();
let info = info.clone();
async move {
if channels.is_empty() {
return Ok(());
}
let (contest, _, _) =
codeforces::Contest::standings(client, rc.contest_id, |f| f.limit(1, 1))
.await?;
channels
.iter()
.map(|channel| {
channel.send_message(http.http(), |e| {
e.content(format!("Rating change for {}!", user_id.mention()))
.embed(|c| {
crate::embed::rating_change_embed(
&rc,
&info,
&contest,
&user_id.mention(),
c,
)
})
})
})
.collect::<stream::FuturesUnordered<_>>()
.map(|v| v.map(|_| ()))
.try_collect::<()>()
.await?;
let r: Result<_> = Ok(());
r
}
})
.collect::<stream::FuturesUnordered<_>>()
.try_collect::<()>()
.await?;
Ok(())
}

199
youmubot-cf/src/hook.rs
View file

@ -1,13 +1,13 @@
use chrono::{TimeZone, Utc};
use codeforces::{Contest, Problem};
use codeforces::{Client, Contest, Problem};
use dashmap::DashMap as HashMap;
use lazy_static::lazy_static;
use rayon::{iter::Either, prelude::*};
use regex::{Captures, Regex};
use serenity::{
builder::CreateEmbed, framework::standard::CommandError, model::channel::Message,
utils::MessageBuilder,
};
use std::{collections::HashMap, sync::Arc};
use std::{sync::Arc, time::Instant};
use youmubot_prelude::*;
lazy_static! {
@ -27,106 +27,132 @@ enum ContestOrProblem {
}
/// Caches the contest list.
#[derive(Clone, Debug, Default)]
pub struct ContestCache(Arc<RwLock<HashMap<u64, (Contest, Option<Vec<Problem>>)>>>);
pub struct ContestCache {
contests: HashMap<u64, (Contest, Option<Vec<Problem>>)>,
all_list: RwLock<(Vec<Contest>, Instant)>,
http: Arc<Client>,
}
impl TypeMapKey for ContestCache {
type Value = ContestCache;
}
impl ContestCache {
fn get(
&self,
http: &<HTTPClient as TypeMapKey>::Value,
contest_id: u64,
) -> Result<(Contest, Option<Vec<Problem>>), CommandError> {
let rl = self.0.read();
match rl.get(&contest_id) {
Some(r @ (_, Some(_))) => Ok(r.clone()),
Some((c, None)) => match Contest::standings(http, contest_id, |f| f.limit(1, 1)) {
Ok((c, p, _)) => Ok({
drop(rl);
let mut v = self.0.write();
let v = v.entry(contest_id).or_insert((c, None));
v.1 = Some(p);
v.clone()
}),
Err(_) => Ok((c.clone(), None)),
},
None => {
drop(rl);
// Step 1: try to fetch it individually
match Contest::standings(http, contest_id, |f| f.limit(1, 1)) {
Ok((c, p, _)) => Ok(self
.0
.write()
.entry(contest_id)
.or_insert((c, Some(p)))
.clone()),
Err(codeforces::Error::Codeforces(s)) if s.ends_with("has not started") => {
// Fetch the entire list
{
let mut m = self.0.write();
let contests = Contest::list(http, contest_id > 100_000)?;
contests.into_iter().for_each(|c| {
m.entry(c.id).or_insert((c, None));
});
}
self.0
.read()
.get(&contest_id)
.cloned()
.ok_or("No contest found".into())
}
Err(e) => Err(e.into()),
}
// Step 2: try to fetch the entire list.
/// Creates a new, empty cache.
pub async fn new(http: Arc<Client>) -> Result<Self> {
let contests_list = Contest::list(&*http, true).await?;
Ok(Self {
contests: HashMap::new(),
all_list: RwLock::new((contests_list, Instant::now())),
http,
})
}
/// Gets a contest from the cache, fetching from upstream if possible.
pub async fn get(&self, contest_id: u64) -> Result<(Contest, Option<Vec<Problem>>)> {
if let Some(v) = self.contests.get(&contest_id) {
if v.1.is_some() {
return Ok(v.clone());
}
}
self.get_and_store_contest(contest_id).await
}
async fn get_and_store_contest(
&self,
contest_id: u64,
) -> Result<(Contest, Option<Vec<Problem>>)> {
let (c, p) = match Contest::standings(&*self.http, contest_id, |f| f.limit(1, 1)).await {
Ok((c, p, _)) => (c, Some(p)),
Err(codeforces::Error::Codeforces(s)) if s.ends_with("has not started") => {
let c = self.get_from_list(contest_id).await?;
(c, None)
}
Err(v) => return Err(Error::from(v)),
};
self.contests.insert(contest_id, (c, p));
Ok(self.contests.get(&contest_id).unwrap().clone())
}
async fn get_from_list(&self, contest_id: u64) -> Result<Contest> {
let last_updated = self.all_list.read().await.1.clone();
if Instant::now() - last_updated > std::time::Duration::from_secs(60 * 60) {
// We update at most once an hour.
*self.all_list.write().await =
(Contest::list(&*self.http, true).await?, Instant::now());
}
self.all_list
.read()
.await
.0
.iter()
.find(|v| v.id == contest_id)
.cloned()
.ok_or_else(|| Error::msg("Contest not found"))
}
}
/// Prints info whenever a problem or contest (or more) is sent on a channel.
pub fn codeforces_info_hook(ctx: &mut Context, m: &Message) {
if m.author.bot {
return;
pub struct InfoHook;
#[async_trait]
impl Hook for InfoHook {
async fn call(&mut self, ctx: &Context, m: &Message) -> Result<()> {
if m.author.bot {
return Ok(());
}
let data = ctx.data.read().await;
let contest_cache = data.get::<ContestCache>().unwrap();
let matches = parse(&m.content[..], contest_cache)
.collect::<Vec<_>>()
.await;
if !matches.is_empty() {
m.channel_id
.send_message(&ctx, |c| {
c.content("Here are the info of the given Codeforces links!")
.embed(|e| print_info_message(&matches[..], e))
})
.await?;
}
Ok(())
}
let http = ctx.data.get_cloned::<HTTPClient>();
let contest_cache = ctx.data.get_cloned::<ContestCache>();
}
fn parse<'a>(
content: &'a str,
contest_cache: &'a ContestCache,
) -> impl stream::Stream<Item = (ContestOrProblem, &'a str)> + 'a {
let matches = CONTEST_LINK
.captures_iter(&m.content)
.chain(PROBLEMSET_LINK.captures_iter(&m.content))
// .collect::<Vec<_>>()
// .into_par_iter()
.filter_map(
|v| match parse_capture(http.clone(), contest_cache.clone(), v) {
Ok(v) => Some(v),
Err(e) => {
dbg!(e);
None
}
},
)
.collect::<Vec<_>>();
if !matches.is_empty() {
m.channel_id
.send_message(&ctx, |c| {
c.content("Here are the info of the given Codeforces links!")
.embed(|e| print_info_message(&matches[..], e))
})
.ok();
}
.captures_iter(content)
.chain(PROBLEMSET_LINK.captures_iter(content))
.map(|v| parse_capture(contest_cache, v))
.collect::<stream::FuturesUnordered<_>>()
.filter_map(|v| future::ready(v.ok()));
matches
}
fn print_info_message<'a>(
info: &[(ContestOrProblem, &str)],
e: &'a mut CreateEmbed,
) -> &'a mut CreateEmbed {
let (mut problems, contests): (Vec<_>, Vec<_>) =
info.par_iter().partition_map(|(v, l)| match v {
ContestOrProblem::Problem(p) => Either::Left((p, l)),
ContestOrProblem::Contest(c, p) => Either::Right((c, p, l)),
});
let (problems, contests): (Vec<_>, Vec<_>) = info.iter().partition(|(v, _)| match v {
ContestOrProblem::Problem(_) => true,
ContestOrProblem::Contest(_, _) => false,
});
let mut problems = problems
.into_iter()
.map(|(v, l)| match v {
ContestOrProblem::Problem(p) => (p, l),
_ => unreachable!(),
})
.collect::<Vec<_>>();
let contests = contests
.into_iter()
.map(|(v, l)| match v {
ContestOrProblem::Contest(c, p) => (c, p, l),
_ => unreachable!(),
})
.collect::<Vec<_>>();
problems.sort_by(|(a, _), (b, _)| a.rating.unwrap_or(1500).cmp(&b.rating.unwrap_or(1500)));
let mut m = MessageBuilder::new();
if !problems.is_empty() {
@ -190,9 +216,8 @@ fn print_info_message<'a>(
e.description(m.build())
}
fn parse_capture<'a>(
http: <HTTPClient as TypeMapKey>::Value,
contest_cache: ContestCache,
async fn parse_capture<'a>(
contest_cache: &ContestCache,
cap: Captures<'a>,
) -> Result<(ContestOrProblem, &'a str), CommandError> {
let contest_id: u64 = cap
@ -200,7 +225,7 @@ fn parse_capture<'a>(
.ok_or(CommandError::from("Contest not captured"))?
.as_str()
.parse()?;
let (contest, problems) = contest_cache.get(&http, contest_id)?;
let (contest, problems) = contest_cache.get(contest_id).await?;
match cap.name("problem") {
Some(p) => {
for problem in problems.ok_or(CommandError::from("Contest hasn't started"))? {

410
youmubot-cf/src/lib.rs
View file

@ -2,12 +2,12 @@ use codeforces::Contest;
use serenity::{
framework::standard::{
macros::{command, group},
Args, CommandError as Error, CommandResult,
Args, CommandResult,
},
model::channel::Message,
utils::MessageBuilder,
};
use std::{collections::HashMap, time::Duration};
use std::{collections::HashMap, sync::Arc, time::Duration};
use youmubot_prelude::*;
mod announcer;
@ -18,16 +18,26 @@ mod hook;
/// Live-commentating a Codeforces round.
mod live;
/// The TypeMapKey holding the Client.
struct CFClient;
impl TypeMapKey for CFClient {
type Value = Arc<codeforces::Client>;
}
use db::{CfSavedUsers, CfUser};
pub use hook::codeforces_info_hook;
pub use hook::InfoHook;
/// Sets up the CF databases.
pub fn setup(path: &std::path::Path, data: &mut ShareMap, announcers: &mut AnnouncerHandler) {
pub async fn setup(path: &std::path::Path, data: &mut TypeMap, announcers: &mut AnnouncerHandler) {
CfSavedUsers::insert_into(data, path.join("cf_saved_users.yaml"))
.expect("Must be able to set up DB");
data.insert::<hook::ContestCache>(hook::ContestCache::default());
announcers.add("codeforces", announcer::updates);
let http = data.get::<HTTPClient>().unwrap();
let client = Arc::new(codeforces::Client::new(http.clone()));
data.insert::<hook::ContestCache>(hook::ContestCache::new(client.clone()).await.unwrap());
data.insert::<CFClient>(client);
announcers.add("codeforces", announcer::Announcer);
}
#[group]
@ -43,40 +53,46 @@ pub struct Codeforces;
#[usage = "[handle or tag = yourself]"]
#[example = "natsukagami"]
#[max_args(1)]
pub fn profile(ctx: &mut Context, m: &Message, mut args: Args) -> CommandResult {
pub async fn profile(ctx: &Context, m: &Message, mut args: Args) -> CommandResult {
let data = ctx.data.read().await;
let handle = args
.single::<UsernameArg>()
.unwrap_or(UsernameArg::mention(m.author.id));
let http = ctx.data.get_cloned::<HTTPClient>();
let http = data.get::<CFClient>().unwrap();
let handle = match handle {
UsernameArg::Raw(s) => s,
UsernameArg::Tagged(u) => {
let db = CfSavedUsers::open(&*ctx.data.read());
let db = db.borrow()?;
match db.get(&u) {
Some(v) => v.handle.clone(),
let db = CfSavedUsers::open(&*data);
let user = db.borrow()?.get(&u).map(|u| u.handle.clone());
match user {
Some(v) => v,
None => {
m.reply(&ctx, "no saved account found.")?;
m.reply(&ctx, "no saved account found.").await?;
return Ok(());
}
}
}
};
let account = codeforces::User::info(&http, &[&handle[..]])?
let account = codeforces::User::info(&http, &[&handle[..]])
.await?
.into_iter()
.next();
match account {
Some(v) => m.channel_id.send_message(&ctx, |send| {
send.content(format!(
"{}: Here is the user that you requested",
m.author.mention()
))
.embed(|e| embed::user_embed(&v, e))
}),
None => m.reply(&ctx, "User not found"),
Some(v) => {
m.channel_id
.send_message(&ctx, |send| {
send.content(format!(
"{}: Here is the user that you requested",
m.author.mention()
))
.embed(|e| embed::user_embed(&v, e))
})
.await
}
None => m.reply(&ctx, "User not found").await,
}?;
Ok(())
@ -86,28 +102,32 @@ pub fn profile(ctx: &mut Context, m: &Message, mut args: Args) -> CommandResult
#[description = "Link your Codeforces account to the Discord account, to enjoy Youmu's tracking capabilities."]
#[usage = "[handle]"]
#[num_args(1)]
pub fn save(ctx: &mut Context, m: &Message, mut args: Args) -> CommandResult {
pub async fn save(ctx: &Context, m: &Message, mut args: Args) -> CommandResult {
let data = ctx.data.read().await;
let handle = args.single::<String>()?;
let http = ctx.data.get_cloned::<HTTPClient>();
let http = data.get::<CFClient>().unwrap();
let account = codeforces::User::info(&http, &[&handle[..]])?
let account = codeforces::User::info(&http, &[&handle[..]])
.await?
.into_iter()
.next();
match account {
None => {
m.reply(&ctx, "cannot find an account with such handle")?;
m.reply(&ctx, "cannot find an account with such handle")
.await?;
}
Some(acc) => {
// Collect rating changes data.
let rating_changes = acc.rating_changes(&http)?;
let db = CfSavedUsers::open(&*ctx.data.read());
let mut db = db.borrow_mut()?;
let rating_changes = acc.rating_changes(&http).await?;
let mut db = CfSavedUsers::open(&*data);
m.reply(
&ctx,
format!("account `{}` has been linked to your account.", &acc.handle),
)?;
db.insert(m.author.id, CfUser::save(acc, rating_changes));
)
.await?;
db.borrow_mut()?
.insert(m.author.id, CfUser::save(acc, rating_changes));
}
}
@ -118,9 +138,10 @@ pub fn save(ctx: &mut Context, m: &Message, mut args: Args) -> CommandResult {
#[description = "See the leaderboard of all people in the server."]
#[only_in(guilds)]
#[num_args(0)]
pub fn ranks(ctx: &mut Context, m: &Message) -> CommandResult {
pub async fn ranks(ctx: &Context, m: &Message) -> CommandResult {
let data = ctx.data.read().await;
let everyone = {
let db = CfSavedUsers::open(&*ctx.data.read());
let db = CfSavedUsers::open(&*data);
let db = db.borrow()?;
db.iter()
.map(|(k, v)| (k.clone(), v.clone()))
@ -129,84 +150,98 @@ pub fn ranks(ctx: &mut Context, m: &Message) -> CommandResult {
let guild = m.guild_id.expect("Guild-only command");
let mut ranks = everyone
.into_iter()
.filter_map(|(id, cf_user)| guild.member(&ctx, id).ok().map(|mem| (mem, cf_user)))
.collect::<Vec<_>>();
.map(|(id, cf_user)| {
guild
.member(&ctx, id)
.map(|mem| mem.map(|mem| (mem, cf_user)))
})
.collect::<stream::FuturesUnordered<_>>()
.filter_map(|v| future::ready(v.ok()))
.collect::<Vec<_>>()
.await;
ranks.sort_by(|(_, a), (_, b)| b.rating.unwrap_or(-1).cmp(&a.rating.unwrap_or(-1)));
if ranks.is_empty() {
m.reply(&ctx, "No saved users in this server.")?;
m.reply(&ctx, "No saved users in this server.").await?;
return Ok(());
}
let ranks = Arc::new(ranks);
const ITEMS_PER_PAGE: usize = 10;
let total_pages = (ranks.len() + ITEMS_PER_PAGE - 1) / ITEMS_PER_PAGE;
let last_updated = ranks.iter().map(|(_, cfu)| cfu.last_update).min().unwrap();
ctx.data.get_cloned::<ReactionWatcher>().paginate_fn(
ctx.clone(),
m.channel_id,
move |page, e| {
let page = page as usize;
let start = ITEMS_PER_PAGE * page;
let end = ranks.len().min(start + ITEMS_PER_PAGE);
if start >= end {
return (e, Err(Error::from("No more pages")));
}
let ranks = &ranks[start..end];
paginate(
move |page, ctx, msg| {
let ranks = ranks.clone();
Box::pin(async move {
let page = page as usize;
let start = ITEMS_PER_PAGE * page;
let end = ranks.len().min(start + ITEMS_PER_PAGE);
if start >= end {
return Ok(false);
}
let ranks = &ranks[start..end];
let handle_width = ranks.iter().map(|(_, cfu)| cfu.handle.len()).max().unwrap();
let username_width = ranks
.iter()
.map(|(mem, _)| mem.distinct().len())
.max()
.unwrap();
let handle_width = ranks.iter().map(|(_, cfu)| cfu.handle.len()).max().unwrap();
let username_width = ranks
.iter()
.map(|(mem, _)| mem.distinct().len())
.max()
.unwrap();
let mut m = MessageBuilder::new();
m.push_line("```");
let mut m = MessageBuilder::new();
m.push_line("```");
// Table header
m.push_line(format!(
"Rank | Rating | {:hw$} | {:uw$}",
"Handle",
"Username",
hw = handle_width,
uw = username_width
));
m.push_line(format!(
"----------------{:->hw$}---{:->uw$}",
"",
"",
hw = handle_width,
uw = username_width
));
for (id, (mem, cfu)) in ranks.iter().enumerate() {
let id = id + start + 1;
// Table header
m.push_line(format!(
"{:>4} | {:>6} | {:hw$} | {:uw$}",
format!("#{}", id),
cfu.rating
.map(|v| v.to_string())
.unwrap_or("----".to_owned()),
cfu.handle,
mem.distinct(),
"Rank | Rating | {:hw$} | {:uw$}",
"Handle",
"Username",
hw = handle_width,
uw = username_width
));
m.push_line(format!(
"----------------{:->hw$}---{:->uw$}",
"",
"",
hw = handle_width,
uw = username_width
));
}
m.push_line("```");
m.push(format!(
"Page **{}/{}**. Last updated **{}**",
page + 1,
total_pages,
last_updated.to_rfc2822()
));
for (id, (mem, cfu)) in ranks.iter().enumerate() {
let id = id + start + 1;
m.push_line(format!(
"{:>4} | {:>6} | {:hw$} | {:uw$}",
format!("#{}", id),
cfu.rating
.map(|v| v.to_string())
.unwrap_or("----".to_owned()),
cfu.handle,
mem.distinct(),
hw = handle_width,
uw = username_width
));
}
(e.content(m.build()), Ok(()))
m.push_line("```");
m.push(format!(
"Page **{}/{}**. Last updated **{}**",
page + 1,
total_pages,
last_updated.to_rfc2822()
));
msg.edit(ctx, |f| f.content(m.build())).await?;
Ok(true)
})
},
ctx,
m.channel_id,
std::time::Duration::from_secs(60),
)?;
)
.await?;
Ok(())
}
@ -216,23 +251,27 @@ pub fn ranks(ctx: &mut Context, m: &Message) -> CommandResult {
#[usage = "[the contest id]"]
#[num_args(1)]
#[only_in(guilds)]
pub fn contestranks(ctx: &mut Context, m: &Message, mut args: Args) -> CommandResult {
pub async fn contestranks(ctx: &Context, m: &Message, mut args: Args) -> CommandResult {
let data = ctx.data.read().await;
let contest_id: u64 = args.single()?;
let guild = m.guild_id.unwrap(); // Guild-only command
let members = CfSavedUsers::open(&*ctx.data.read()).borrow()?.clone();
let members = CfSavedUsers::open(&*data).borrow()?.clone();
let members = members
.into_iter()
.filter_map(|(user_id, cf_user)| {
.map(|(user_id, cf_user)| {
guild
.member(&ctx, user_id)
.ok()
.map(|v| (cf_user.handle, v))
.map(|v| v.map(|v| (cf_user.handle, v)))
})
.collect::<HashMap<_, _>>();
let http = ctx.data.get_cloned::<HTTPClient>();
let (contest, problems, ranks) = Contest::standings(&http, contest_id, |f| {
.collect::<stream::FuturesUnordered<_>>()
.filter_map(|v| future::ready(v.ok()))
.collect::<HashMap<_, _>>()
.await;
let http = data.get::<CFClient>().unwrap();
let (contest, problems, ranks) = Contest::standings(http, contest_id, |f| {
f.handles(members.iter().map(|(k, _)| k.clone()).collect())
})?;
})
.await?;
// Table me
let ranks = ranks
@ -252,100 +291,111 @@ pub fn contestranks(ctx: &mut Context, m: &Message, mut args: Args) -> CommandRe
.collect::<Vec<_>>();
if ranks.is_empty() {
m.reply(&ctx, "No one in this server participated in the contest...")?;
m.reply(&ctx, "No one in this server participated in the contest...")
.await?;
return Ok(());
}
let ranks = Arc::new(ranks);
const ITEMS_PER_PAGE: usize = 10;
let total_pages = (ranks.len() + ITEMS_PER_PAGE - 1) / ITEMS_PER_PAGE;
ctx.data.get_cloned::<ReactionWatcher>().paginate_fn(
ctx.clone(),
m.channel_id,
move |page, e| {
let page = page as usize;
let start = page * ITEMS_PER_PAGE;
let end = ranks.len().min(start + ITEMS_PER_PAGE);
if start >= end {
return (e, Err(Error::from("no more pages to show")));
}
let ranks = &ranks[start..end];
let hw = ranks
.iter()
.map(|(mem, handle, _)| format!("{} ({})", handle, mem.distinct()).len())
.max()
.unwrap_or(0)
.max(6);
let hackw = ranks
.iter()
.map(|(_, _, row)| {
format!(
"{}/{}",
row.successful_hack_count, row.unsuccessful_hack_count
)
.len()
})
.max()
.unwrap_or(0)
.max(5);
paginate(
move |page, ctx, msg| {
let contest = contest.clone();
let problems = problems.clone();
let ranks = ranks.clone();
Box::pin(async move {
let page = page as usize;
let start = page * ITEMS_PER_PAGE;
let end = ranks.len().min(start + ITEMS_PER_PAGE);
if start >= end {
return Ok(false);
}
let ranks = &ranks[start..end];
let hw = ranks
.iter()
.map(|(mem, handle, _)| format!("{} ({})", handle, mem.distinct()).len())
.max()
.unwrap_or(0)
.max(6);
let hackw = ranks
.iter()
.map(|(_, _, row)| {
format!(
"{}/{}",
row.successful_hack_count, row.unsuccessful_hack_count
)
.len()
})
.max()
.unwrap_or(0)
.max(5);
let mut table = MessageBuilder::new();
let mut header = MessageBuilder::new();
// Header
header.push(format!(
" Rank | {:hw$} | Total | {:hackw$}",
"Handle",
"Hacks",
hw = hw,
hackw = hackw
));
for p in &problems {
header.push(format!(" | {:4}", p.index));
}
let header = header.build();
table
.push_line(&header)
.push_line(format!("{:-<w$}", "", w = header.len()));
// Body
for (mem, handle, row) in ranks {
table.push(format!(
"{:>5} | {:<hw$} | {:>5.0} | {:<hackw$}",
row.rank,
format!("{} ({})", handle, mem.distinct()),
row.points,
format!(
"{}/{}",
row.successful_hack_count, row.unsuccessful_hack_count
),
let mut table = MessageBuilder::new();
let mut header = MessageBuilder::new();
// Header
header.push(format!(
" Rank | {:hw$} | Total | {:hackw$}",
"Handle",
"Hacks",
hw = hw,
hackw = hackw
));
for p in &row.problem_results {
table.push(" | ");
if p.points > 0.0 {
table.push(format!("{:^4.0}", p.points));
} else if let Some(_) = p.best_submission_time_seconds {
table.push(format!("{:^4}", "?"));
} else if p.rejected_attempt_count > 0 {
table.push(format!("{:^4}", format!("-{}", p.rejected_attempt_count)));
} else {
table.push(format!("{:^4}", ""));
}
for p in &problems {
header.push(format!(" | {:4}", p.index));
}
table.push_line("");
}
let header = header.build();
table
.push_line(&header)
.push_line(format!("{:-<w$}", "", w = header.len()));
let mut m = MessageBuilder::new();
m.push_bold_safe(&contest.name)
.push(" ")
.push_line(contest.url())
.push_codeblock(table.build(), None)
.push_line(format!("Page **{}/{}**", page + 1, total_pages));
(e.content(m.build()), Ok(()))
// Body
for (mem, handle, row) in ranks {
table.push(format!(
"{:>5} | {:<hw$} | {:>5.0} | {:<hackw$}",
row.rank,
format!("{} ({})", handle, mem.distinct()),
row.points,
format!(
"{}/{}",
row.successful_hack_count, row.unsuccessful_hack_count
),
hw = hw,
hackw = hackw
));
for p in &row.problem_results {
table.push(" | ");
if p.points > 0.0 {
table.push(format!("{:^4.0}", p.points));
} else if let Some(_) = p.best_submission_time_seconds {
table.push(format!("{:^4}", "?"));
} else if p.rejected_attempt_count > 0 {
table.push(format!("{:^4}", format!("-{}", p.rejected_attempt_count)));
} else {
table.push(format!("{:^4}", ""));
}
}
table.push_line("");
}
let mut m = MessageBuilder::new();
m.push_bold_safe(&contest.name)
.push(" ")
.push_line(contest.url())
.push_codeblock(table.build(), None)
.push_line(format!("Page **{}/{}**", page + 1, total_pages));
msg.edit(ctx, |e| e.content(m.build())).await?;
Ok(true)
})
},
ctx,
m.channel_id,
Duration::from_secs(60),
)
.await?;
Ok(())
}
#[command]
@ -354,10 +404,10 @@ pub fn contestranks(ctx: &mut Context, m: &Message, mut args: Args) -> CommandRe
#[num_args(1)]
#[required_permissions(MANAGE_CHANNELS)]
#[only_in(guilds)]
pub fn watch(ctx: &mut Context, m: &Message, mut args: Args) -> CommandResult {
pub async fn watch(ctx: &Context, m: &Message, mut args: Args) -> CommandResult {
let contest_id: u64 = args.single()?;
live::watch_contest(ctx, m.guild_id.unwrap(), m.channel_id, contest_id)?;
live::watch_contest(ctx, m.guild_id.unwrap(), m.channel_id, contest_id).await?;
Ok(())
}

102
youmubot-cf/src/live.rs
View file

@ -1,8 +1,6 @@
use crate::db::CfSavedUsers;
use crate::{db::CfSavedUsers, CFClient};
use codeforces::{Contest, ContestPhase, Problem, ProblemResult, ProblemResultType, RanklistRow};
use rayon::prelude::*;
use serenity::{
framework::standard::{CommandError, CommandResult},
model::{
guild::Member,
id::{ChannelId, GuildId, UserId},
@ -21,56 +19,64 @@ struct MemberResult {
/// Watch and commentate a contest.
///
/// Does the thing on a channel, block until the contest ends.
pub fn watch_contest(
ctx: &mut Context,
pub async fn watch_contest(
ctx: &Context,
guild: GuildId,
channel: ChannelId,
contest_id: u64,
) -> CommandResult {
let db = CfSavedUsers::open(&*ctx.data.read()).borrow()?.clone();
) -> Result<()> {
let data = ctx.data.read().await;
let db = CfSavedUsers::open(&*data).borrow()?.clone();
let http = ctx.http.clone();
// Collect an initial member list.
// This never changes during the scan.
let mut member_results: HashMap<UserId, MemberResult> = db
.into_par_iter()
.filter_map(|(user_id, cfu)| {
let member = guild.member(http.clone().as_ref(), user_id).ok();
match member {
Some(m) => Some((
user_id,
MemberResult {
member: m,
handle: cfu.handle,
row: None,
},
)),
None => None,
.into_iter()
.map(|(user_id, cfu)| {
let http = http.clone();
async move {
guild.member(http, user_id).await.map(|m| {
(
user_id,
MemberResult {
member: m,
handle: cfu.handle,
row: None,
},
)
})
}
})
.collect();
.collect::<stream::FuturesUnordered<_>>()
.filter_map(|v| future::ready(v.ok()))
.collect()
.await;
let http = ctx.data.get_cloned::<HTTPClient>();
let (mut contest, _, _) = Contest::standings(&http, contest_id, |f| f.limit(1, 1))?;
let http = data.get::<CFClient>().unwrap();
let (mut contest, _, _) = Contest::standings(&http, contest_id, |f| f.limit(1, 1)).await?;
channel.send_message(&ctx, |e| {
e.content(format!(
"Youmu is watching contest **{}**, with the following members:\n{}",
contest.name,
member_results
.iter()
.map(|(_, m)| format!("- {} as **{}**", m.member.distinct(), m.handle))
.collect::<Vec<_>>()
.join("\n"),
))
})?;
channel
.send_message(&ctx, |e| {
e.content(format!(
"Youmu is watching contest **{}**, with the following members:\n{}",
contest.name,
member_results
.iter()
.map(|(_, m)| format!("- {} as **{}**", m.member.distinct(), m.handle))
.collect::<Vec<_>>()
.join("\n"),
))
})
.await?;
loop {
if let Ok(messages) = scan_changes(http.clone(), &mut member_results, &mut contest) {
if let Ok(messages) = scan_changes(&*http, &mut member_results, &mut contest).await {
for message in messages {
channel
.send_message(&ctx, |e| {
e.content(format!("**{}**: {}", contest.name, message))
})
.await
.ok();
}
}
@ -78,7 +84,7 @@ pub fn watch_contest(
break;
}
// Sleep for a minute
std::thread::sleep(std::time::Duration::from_secs(60));
tokio::time::delay_for(std::time::Duration::from_secs(60)).await;
}
// Announce the final results
@ -93,12 +99,14 @@ pub fn watch_contest(
ranks.sort_by(|(_, a), (_, b)| a.rank.cmp(&b.rank));
if ranks.is_empty() {
channel.send_message(&ctx, |e| {
e.content(format!(
"**{}** has ended, but I can't find anyone in this server on the scoreboard...",
contest.name
))
})?;
channel
.send_message(&ctx, |e| {
e.content(format!(
"**{}** has ended, but I can't find anyone in this server on the scoreboard...",
contest.name
))
})
.await?;
return Ok(());
}
@ -115,23 +123,23 @@ pub fn watch_contest(
row.problem_results.iter().map(|p| format!("{:.0}", p.points)).collect::<Vec<_>>().join("/"),
row.successful_hack_count,
row.unsuccessful_hack_count,
)).collect::<Vec<_>>().join("\n"))))?;
)).collect::<Vec<_>>().join("\n")))).await?;
Ok(())
}
fn scan_changes(
http: <HTTPClient as TypeMapKey>::Value,
async fn scan_changes(
http: &codeforces::Client,
members: &mut HashMap<UserId, MemberResult>,
contest: &mut Contest,
) -> Result<Vec<String>, CommandError> {
) -> Result<Vec<String>> {
let mut messages: Vec<String> = vec![];
let (updated_contest, problems, ranks) = {
let handles = members
.iter()
.map(|(_, h)| h.handle.clone())
.collect::<Vec<_>>();
Contest::standings(&http, contest.id, |f| f.handles(handles))?
Contest::standings(&http, contest.id, |f| f.handles(handles)).await?
};
// Change of phase.
if contest.phase != updated_contest.phase {

4
youmubot-core/Cargo.toml
View file

@ -7,11 +7,13 @@ edition = "2018"
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
[dependencies]
serenity = "0.8"
serenity = { version = "0.9.0-rc.0", features = ["collector"] }
rand = "0.7"
serde = { version = "1", features = ["derive"] }
chrono = "0.4"
static_assertions = "1.1"
futures-util = "0.3"
tokio = { version = "0.2", features = ["time"] }
youmubot-db = { path = "../youmubot-db" }
youmubot-prelude = { path = "../youmubot-prelude" }

78
youmubot-core/src/admin/mod.rs
View file

@ -1,3 +1,4 @@
use futures_util::{stream, TryStreamExt};
use serenity::{
framework::standard::{
macros::{command, group},
@ -9,7 +10,6 @@ use serenity::{
},
};
use soft_ban::{SOFT_BAN_COMMAND, SOFT_BAN_INIT_COMMAND};
use std::{thread::sleep, time::Duration};
use youmubot_prelude::*;
mod soft_ban;
@ -27,29 +27,34 @@ struct Admin;
#[usage = "clean 50"]
#[min_args(0)]
#[max_args(1)]
fn clean(ctx: &mut Context, msg: &Message, mut args: Args) -> CommandResult {
async fn clean(ctx: &Context, msg: &Message, mut args: Args) -> CommandResult {
let limit = args.single().unwrap_or(10);
let messages = msg
.channel_id
.messages(&ctx.http, |b| b.before(msg.id).limit(limit))?;
let channel = msg.channel_id.to_channel(&ctx)?;
.messages(&ctx.http, |b| b.before(msg.id).limit(limit))
.await?;
let channel = msg.channel_id.to_channel(&ctx).await?;
match &channel {
Channel::Private(_) | Channel::Group(_) => {
let self_id = ctx.http.get_current_application_info()?.id;
Channel::Private(_) => {
let self_id = ctx.http.get_current_application_info().await?.id;
messages
.into_iter()
.filter(|v| v.author.id == self_id)
.try_for_each(|m| m.delete(&ctx))?;
.map(|m| async move { m.delete(&ctx).await })
.collect::<stream::FuturesUnordered<_>>()
.try_collect::<()>()
.await?;
}
_ => {
msg.channel_id
.delete_messages(&ctx.http, messages.into_iter())?;
.delete_messages(&ctx.http, messages.into_iter())
.await?;
}
};
msg.react(&ctx, "🌋")?;
msg.react(&ctx, '🌋').await?;
if let Channel::Guild(_) = &channel {
sleep(Duration::from_secs(2));
msg.delete(&ctx)?;
tokio::time::delay_for(std::time::Duration::from_secs(2)).await;
msg.delete(&ctx).await.ok();
}
Ok(())
@ -58,25 +63,36 @@ fn clean(ctx: &mut Context, msg: &Message, mut args: Args) -> CommandResult {
#[command]
#[required_permissions(ADMINISTRATOR)]
#[description = "Ban an user with a certain reason."]
#[usage = "@user#1234/spam"]
#[usage = "tag user/[reason = none]/[days of messages to delete = 0]"]
#[min_args(1)]
#[max_args(2)]
#[only_in("guilds")]
fn ban(ctx: &mut Context, msg: &Message, mut args: Args) -> CommandResult {
let user = args.single::<UserId>()?.to_user(&ctx)?;
let reason = args
.remains()
.map(|v| format!("`{}`", v))
.unwrap_or("no provided reason".to_owned());
async fn ban(ctx: &Context, msg: &Message, mut args: Args) -> CommandResult {
let user = args.single::<UserId>()?.to_user(&ctx).await?;
let reason = args.single::<String>().map(|v| format!("`{}`", v)).ok();
let dmds = args.single::<u8>().unwrap_or(0);
msg.reply(
&ctx,
format!("🔨 Banning user {} for reason `{}`.", user.tag(), reason),
)?;
msg.guild_id
.ok_or("Can't get guild from message?")? // we had a contract
.ban(&ctx.http, user, &reason)?;
match reason {
Some(reason) => {
msg.reply(
&ctx,
format!("🔨 Banning user {} for reason `{}`.", user.tag(), reason),
)
.await?;
msg.guild_id
.ok_or(Error::msg("Can't get guild from message?"))? // we had a contract
.ban_with_reason(&ctx.http, user, dmds, &reason)
.await?;
}
None => {
msg.reply(&ctx, format!("🔨 Banning user {}.", user.tag()))
.await?;
msg.guild_id
.ok_or(Error::msg("Can't get guild from message?"))? // we had a contract
.ban(&ctx.http, user, dmds)
.await?;
}
}
Ok(())
}
@ -87,14 +103,16 @@ fn ban(ctx: &mut Context, msg: &Message, mut args: Args) -> CommandResult {
#[usage = "@user#1234"]
#[num_args(1)]
#[only_in("guilds")]
fn kick(ctx: &mut Context, msg: &Message, mut args: Args) -> CommandResult {
let user = args.single::<UserId>()?.to_user(&ctx)?;
async fn kick(ctx: &Context, msg: &Message, mut args: Args) -> CommandResult {
let user = args.single::<UserId>()?.to_user(&ctx).await?;
msg.reply(&ctx, format!("🔫 Kicking user {}.", user.tag()))?;
msg.reply(&ctx, format!("🔫 Kicking user {}.", user.tag()))
.await?;
msg.guild_id
.ok_or("Can't get guild from message?")? // we had a contract
.kick(&ctx.http, user)?;
.kick(&ctx.http, user)
.await?;
Ok(())
}

235
youmubot-core/src/admin/soft_ban.rs
View file

@ -1,13 +1,15 @@
use crate::db::{ServerSoftBans, SoftBans};
use chrono::offset::Utc;
use futures_util::{stream, TryStreamExt};
use serenity::{
framework::standard::{macros::command, Args, CommandError as Error, CommandResult},
framework::standard::{macros::command, Args, CommandResult},
model::{
channel::Message,
id::{RoleId, UserId},
id::{GuildId, RoleId, UserId},
},
CacheAndHttp,
};
use std::cmp::max;
use std::sync::Arc;
use youmubot_prelude::*;
#[command]
@ -18,57 +20,56 @@ use youmubot_prelude::*;
#[min_args(1)]
#[max_args(2)]
#[only_in("guilds")]
pub fn soft_ban(ctx: &mut Context, msg: &Message, mut args: Args) -> CommandResult {
let user = args.single::<UserId>()?.to_user(&ctx)?;
pub async fn soft_ban(ctx: &Context, msg: &Message, mut args: Args) -> CommandResult {
let user = args.single::<UserId>()?.to_user(&ctx).await?;
let data = ctx.data.read().await;
let duration = if args.is_empty() {
None
} else {
Some(
args.single::<args::Duration>()
.map_err(|e| Error::from(&format!("{:?}", e)))?,
)
Some(args.single::<args::Duration>()?)
};
let guild = msg.guild_id.ok_or(Error::from("Command is guild only"))?;
let guild = msg.guild_id.ok_or(Error::msg("Command is guild only"))?;
let db = SoftBans::open(&*ctx.data.read());
let mut db = db.borrow_mut()?;
let mut server_ban = db.get_mut(&guild).and_then(|v| match v {
ServerSoftBans::Unimplemented => None,
ServerSoftBans::Implemented(ref mut v) => Some(v),
});
match server_ban {
let mut db = SoftBans::open(&*data);
let val = db
.borrow()?
.get(&guild)
.map(|v| (v.role, v.periodical_bans.get(&user.id).cloned()));
let (role, current_ban_deadline) = match val {
None => {
println!("get here");
msg.reply(&ctx, format!("⚠ This server has not enabled the soft-ban feature. Check out `y!a soft-ban-init`."))?;
msg.reply(&ctx, format!("⚠ This server has not enabled the soft-ban feature. Check out `y!a soft-ban-init`.")).await?;
return Ok(());
}
Some(ref mut server_ban) => {
let mut member = guild.member(&ctx, &user)?;
match duration {
None if member.roles.contains(&server_ban.role) => {
msg.reply(&ctx, format!("⛓ Lifting soft-ban for user {}.", user.tag()))?;
member.remove_role(&ctx, server_ban.role)?;
return Ok(());
}
None => {
msg.reply(&ctx, format!("⛓ Soft-banning user {}.", user.tag()))?;
}
Some(v) => {
let until = Utc::now() + chrono::Duration::from_std(v.0)?;
let until = server_ban
.periodical_bans
.entry(user.id)
.and_modify(|v| *v = max(*v, until))
.or_insert(until);
msg.reply(
&ctx,
format!("⛓ Soft-banning user {} until {}.", user.tag(), until),
)?;
}
}
member.add_role(&ctx, server_ban.role)?;
Some(v) => v,
};
let mut member = guild.member(&ctx, &user).await?;
match duration {
None if member.roles.contains(&role) => {
msg.reply(&ctx, format!("⛓ Lifting soft-ban for user {}.", user.tag()))
.await?;
member.remove_role(&ctx, role).await?;
return Ok(());
}
None => {
msg.reply(&ctx, format!("⛓ Soft-banning user {}.", user.tag()))
.await?;
}
Some(v) => {
// Add the duration into the ban timeout.
let until =
current_ban_deadline.unwrap_or(Utc::now()) + chrono::Duration::from_std(v.0)?;
msg.reply(
&ctx,
format!("⛓ Soft-banning user {} until {}.", user.tag(), until),
)
.await?;
db.borrow_mut()?
.get_mut(&guild)
.map(|v| v.periodical_bans.insert(user.id, until));
}
}
member.add_role(&ctx, role).await?;
Ok(())
}
@ -79,86 +80,90 @@ pub fn soft_ban(ctx: &mut Context, msg: &Message, mut args: Args) -> CommandResu
#[usage = "{soft_ban_role_id}"]
#[num_args(1)]
#[only_in("guilds")]
pub fn soft_ban_init(ctx: &mut Context, msg: &Message, mut args: Args) -> CommandResult {
pub async fn soft_ban_init(ctx: &Context, msg: &Message, mut args: Args) -> CommandResult {
let role_id = args.single::<RoleId>()?;
let guild = msg.guild(&ctx).ok_or(Error::from("Guild-only command"))?;
let guild = guild.read();
let data = ctx.data.read().await;
let guild = msg.guild(&ctx).await.unwrap();
// Check whether the role_id is the one we wanted
if !guild.roles.contains_key(&role_id) {
return Err(Error::from(format!(
Err(Error::msg(format!(
"{} is not a role in this server.",
role_id
)));
)))?;
}
// Check if we already set up
let db = SoftBans::open(&*ctx.data.read());
let mut db = db.borrow_mut()?;
let server = db
.get(&guild.id)
.map(|v| match v {
ServerSoftBans::Unimplemented => false,
_ => true,
})
.unwrap_or(false);
let mut db = SoftBans::open(&*data);
let set_up = db.borrow()?.contains_key(&guild.id);
if !server {
db.insert(guild.id, ServerSoftBans::new_implemented(role_id));
msg.react(&ctx, "👌")?;
Ok(())
if !set_up {
db.borrow_mut()?
.insert(guild.id, ServerSoftBans::new(role_id));
msg.react(&ctx, '👌').await?;
} else {
Err(Error::from("Server already set up soft-bans."))
Err(Error::msg("Server already set up soft-bans."))?
}
Ok(())
}
// Watch the soft bans. Blocks forever.
pub async fn watch_soft_bans(cache_http: Arc<CacheAndHttp>, data: AppData) {
loop {
// Scope so that locks are released
{
// Poll the data for any changes.
let db = data.read().await;
let db = SoftBans::open(&*db);
let mut db = db.borrow().unwrap().clone();
let now = Utc::now();
for (server_id, bans) in db.iter_mut() {
let server_name: String = match server_id.to_partial_guild(&*cache_http.http).await
{
Err(_) => continue,
Ok(v) => v.name,
};
let to_remove: Vec<_> = bans
.periodical_bans
.iter()
.filter_map(|(user, time)| if time <= &now { Some(user) } else { None })
.cloned()
.collect();
if let Err(e) = to_remove
.into_iter()
.map(|user_id| {
bans.periodical_bans.remove(&user_id);
lift_soft_ban_for(
&*cache_http,
*server_id,
&server_name[..],
bans.role,
user_id,
)
})
.collect::<stream::FuturesUnordered<_>>()
.try_collect::<()>()
.await
{
eprintln!("Error while scanning soft-bans list: {}", e)
}
}
}
// Sleep the thread for a minute
tokio::time::delay_for(std::time::Duration::from_secs(60)).await
}
}
// Watch the soft bans.
pub fn watch_soft_bans(client: &serenity::Client) -> impl FnOnce() -> () + 'static {
let cache_http = {
let cache_http = client.cache_and_http.clone();
let cache: serenity::cache::CacheRwLock = cache_http.cache.clone().into();
(cache, cache_http.http.clone())
};
let data = client.data.clone();
return move || {
let cache_http = (&cache_http.0, &*cache_http.1);
loop {
// Scope so that locks are released
{
// Poll the data for any changes.
let db = data.read();
let db = SoftBans::open(&*db);
let mut db = db.borrow_mut().expect("Borrowable");
let now = Utc::now();
for (server_id, soft_bans) in db.iter_mut() {
let server_name: String = match server_id.to_partial_guild(cache_http) {
Err(_) => continue,
Ok(v) => v.name,
};
if let ServerSoftBans::Implemented(ref mut bans) = soft_bans {
let to_remove: Vec<_> = bans
.periodical_bans
.iter()
.filter_map(|(user, time)| if time <= &now { Some(user) } else { None })
.cloned()
.collect();
for user_id in to_remove {
server_id
.member(cache_http, user_id)
.and_then(|mut m| {
println!(
"Soft-ban for `{}` in server `{}` unlifted.",
m.user.read().name,
server_name
);
m.remove_role(cache_http, bans.role)
})
.unwrap_or(());
bans.periodical_bans.remove(&user_id);
}
}
}
}
// Sleep the thread for a minute
std::thread::sleep(std::time::Duration::from_secs(60))
}
};
async fn lift_soft_ban_for(
cache_http: &CacheAndHttp,
server_id: GuildId,
server_name: &str,
ban_role: RoleId,
user_id: UserId,
) -> Result<()> {
let mut m = server_id.member(cache_http, user_id).await?;
println!(
"Soft-ban for `{}` in server `{}` unlifted.",
m.user.name, server_name
);
m.remove_role(&cache_http.http, ban_role).await?;
Ok(())
}

74
youmubot-core/src/community/mod.rs
View file

@ -9,6 +9,7 @@ use serenity::{
},
model::{
channel::{Channel, Message},
id::RoleId,
user::OnlineStatus,
},
utils::MessageBuilder,
@ -30,11 +31,12 @@ struct Community;
#[command]
#[description = r"👑 Randomly choose an active member and mention them!
Note that only online/idle users in the channel are chosen from."]
#[usage = "[title = the chosen one]"]
#[usage = "[limited roles = everyone online] / [title = the chosen one]"]
#[example = "the strongest in Gensokyo"]
#[bucket = "community"]
#[max_args(1)]
pub fn choose(ctx: &mut Context, m: &Message, mut args: Args) -> CommandResult {
#[max_args(2)]
pub async fn choose(ctx: &Context, m: &Message, mut args: Args) -> CommandResult {
let role = args.find::<RoleId>().ok();
let title = if args.is_empty() {
"the chosen one".to_owned()
} else {
@ -42,29 +44,39 @@ pub fn choose(ctx: &mut Context, m: &Message, mut args: Args) -> CommandResult {
};
let users: Result<Vec<_>, Error> = {
let guild = m.guild(&ctx).unwrap();
let guild = guild.read();
let guild = m.guild(&ctx).await.unwrap();
let presences = &guild.presences;
let channel = m.channel_id.to_channel(&ctx)?;
let channel = m.channel_id.to_channel(&ctx).await?;
if let Channel::Guild(channel) = channel {
let channel = channel.read();
Ok(channel
.members(&ctx)?
.members(&ctx)
.await?
.into_iter()
.filter(|v| !v.user.read().bot)
.map(|v| v.user_id())
.filter(|v| !v.user.bot) // Filter out bots
.filter(|v| {
// Filter out only online people
presences
.get(v)
.get(&v.user.id)
.map(|presence| {
presence.status == OnlineStatus::Online
|| presence.status == OnlineStatus::Idle
})
.unwrap_or(false)
})
.collect())
.map(|mem| future::ready(mem))
.collect::<stream::FuturesUnordered<_>>()
.filter_map(|member| async move {
// Filter by role if provided
match role {
Some(role) if member.roles.iter().any(|r| role == *r) => Some(member),
None => Some(member),
_ => None,
}
})
.collect()
.await)
} else {
panic!()
unreachable!()
}
};
let users = users?;
@ -73,7 +85,8 @@ pub fn choose(ctx: &mut Context, m: &Message, mut args: Args) -> CommandResult {
m.reply(
&ctx,
"🍰 Have this cake for yourself because no-one is here for the gods to pick.",
)?;
)
.await?;
return Ok(());
}
@ -83,19 +96,26 @@ pub fn choose(ctx: &mut Context, m: &Message, mut args: Args) -> CommandResult {
&users[uniform.sample(&mut rng)]
};
m.channel_id.send_message(&ctx, |c| {
c.content(
MessageBuilder::new()
.push("👑 The Gensokyo gods have gathered around and decided, out of ")
.push_bold(format!("{}", users.len()))
.push(" potential prayers, ")
.push(winner.mention())
.push(" will be ")
.push_bold_safe(title)
.push(". Congrats! 🎉 🎊 🥳")
.build(),
)
})?;
m.channel_id
.send_message(&ctx, |c| {
c.content(
MessageBuilder::new()
.push("👑 The Gensokyo gods have gathered around and decided, out of ")
.push_bold(format!("{}", users.len()))
.push(" ")
.push(
role.map(|r| r.mention() + "s")
.unwrap_or("potential prayers".to_owned()),
)
.push(", ")
.push(winner.mention())
.push(" will be ")
.push_bold_safe(title)
.push(". Congrats! 🎉 🎊 🥳")
.build(),
)
})
.await?;
Ok(())
}

192
youmubot-core/src/community/roles.rs
View file

@ -1,6 +1,6 @@
use crate::db::Roles as DB;
use serenity::{
framework::standard::{macros::command, Args, CommandError as Error, CommandResult},
framework::standard::{macros::command, Args, CommandResult},
model::{channel::Message, guild::Role, id::RoleId},
utils::MessageBuilder,
};
@ -10,18 +10,22 @@ use youmubot_prelude::*;
#[description = "List all available roles in the server."]
#[num_args(0)]
#[only_in(guilds)]
fn list(ctx: &mut Context, m: &Message, _: Args) -> CommandResult {
async fn list(ctx: &Context, m: &Message, _: Args) -> CommandResult {
let guild_id = m.guild_id.unwrap(); // only_in(guilds)
let data = ctx.data.read().await;
let db = DB::open(&*ctx.data.read());
let db = db.borrow()?;
let roles = db.get(&guild_id).filter(|v| !v.is_empty()).cloned();
let db = DB::open(&*data);
let roles = db
.borrow()?
.get(&guild_id)
.filter(|v| !v.is_empty())
.cloned();
match roles {
None => {
m.reply(&ctx, "No roles available for assigning.")?;
m.reply(&ctx, "No roles available for assigning.").await?;
}
Some(v) => {
let roles = guild_id.to_partial_guild(&ctx)?.roles;
let roles = guild_id.to_partial_guild(&ctx).await?.roles;
let roles: Vec<_> = v
.into_iter()
.filter_map(|(_, role)| roles.get(&role.id).cloned().map(|r| (r, role.description)))
@ -29,108 +33,116 @@ fn list(ctx: &mut Context, m: &Message, _: Args) -> CommandResult {
const ROLES_PER_PAGE: usize = 8;
let pages = (roles.len() + ROLES_PER_PAGE - 1) / ROLES_PER_PAGE;
let watcher = ctx.data.get_cloned::<ReactionWatcher>();
watcher.paginate_fn(
ctx.clone(),
m.channel_id,
move |page, e| {
let page = page as usize;
let start = page * ROLES_PER_PAGE;
let end = roles.len().min(start + ROLES_PER_PAGE);
if end <= start {
return (e, Err(Error::from("No more roles to display")));
}
let roles = &roles[start..end];
let nw = roles // name width
.iter()
.map(|(r, _)| r.name.len())
.max()
.unwrap()
.max(6);
let idw = roles[0].0.id.to_string().len();
let dw = roles
.iter()
.map(|v| v.1.len())
.max()
.unwrap()
.max(" Description ".len());
let mut m = MessageBuilder::new();
m.push_line("```");
paginate(
|page, ctx, msg| {
let roles = roles.clone();
Box::pin(async move {
let page = page as usize;
let start = page * ROLES_PER_PAGE;
let end = roles.len().min(start + ROLES_PER_PAGE);
if end <= start {
return Ok(false);
}
let roles = &roles[start..end];
let nw = roles // name width
.iter()
.map(|(r, _)| r.name.len())
.max()
.unwrap()
.max(6);
let idw = roles[0].0.id.to_string().len();
let dw = roles
.iter()
.map(|v| v.1.len())
.max()
.unwrap()
.max(" Description ".len());
let mut m = MessageBuilder::new();
m.push_line("```");
// Table header
m.push_line(format!(
"{:nw$} | {:idw$} | {:dw$}",
"Name",
"ID",
"Description",
nw = nw,
idw = idw,
dw = dw,
));
m.push_line(format!(
"{:->nw$}---{:->idw$}---{:->dw$}",
"",
"",
"",
nw = nw,
idw = idw,
dw = dw,
));
for (role, description) in roles.iter() {
// Table header
m.push_line(format!(
"{:nw$} | {:idw$} | {:dw$}",
role.name,
role.id,
description,
"Name",
"ID",
"Description",
nw = nw,
idw = idw,
dw = dw,
));
m.push_line(format!(
"{:->nw$}---{:->idw$}---{:->dw$}",
"",
"",
"",
nw = nw,
idw = idw,
dw = dw,
));
}
m.push_line("```");
m.push(format!("Page **{}/{}**", page + 1, pages));
(e.content(m.build()), Ok(()))
for (role, description) in roles.iter() {
m.push_line(format!(
"{:nw$} | {:idw$} | {:dw$}",
role.name,
role.id,
description,
nw = nw,
idw = idw,
dw = dw,
));
}
m.push_line("```");
m.push(format!("Page **{}/{}**", page + 1, pages));
msg.edit(ctx, |f| f.content(m.to_string())).await?;
Ok(true)
})
},
ctx,
m.channel_id,
std::time::Duration::from_secs(60 * 10),
)?;
)
.await?;
}
};
Ok(())
}
// async fn list_pager(
#[command("role")]
#[description = "Toggle a role by its name or ID."]
#[example = "\"IELTS / TOEFL\""]
#[num_args(1)]
#[only_in(guilds)]
fn toggle(ctx: &mut Context, m: &Message, mut args: Args) -> CommandResult {
async fn toggle(ctx: &Context, m: &Message, mut args: Args) -> CommandResult {
let role = args.single_quoted::<String>()?;
let guild_id = m.guild_id.unwrap();
let roles = guild_id.to_partial_guild(&ctx)?.roles;
let role = role_from_string(&role, &roles);
let guild = guild_id.to_partial_guild(&ctx).await?;
let role = role_from_string(&role, &guild.roles);
match role {
None => {
m.reply(&ctx, "No such role exists")?;
m.reply(&ctx, "No such role exists").await?;
}
Some(role)
if !DB::open(&*ctx.data.read())
if !DB::open(&*ctx.data.read().await)
.borrow()?
.get(&guild_id)
.map(|g| g.contains_key(&role.id))
.unwrap_or(false) =>
{
m.reply(&ctx, "This role is not self-assignable. Check the `listroles` command to see which role can be assigned.")?;
m.reply(&ctx, "This role is not self-assignable. Check the `listroles` command to see which role can be assigned.").await?;
}
Some(role) => {
let mut member = m.member(&ctx).ok_or(Error::from("Cannot find member"))?;
let mut member = guild.member(&ctx, m.author.id).await.unwrap();
if member.roles.contains(&role.id) {
member.remove_role(&ctx, &role)?;
m.reply(&ctx, format!("Role `{}` has been removed.", role.name))?;
member.remove_role(&ctx, &role).await?;
m.reply(&ctx, format!("Role `{}` has been removed.", role.name))
.await?;
} else {
member.add_role(&ctx, &role)?;
m.reply(&ctx, format!("Role `{}` has been assigned.", role.name))?;
member.add_role(&ctx, &role).await?;
m.reply(&ctx, format!("Role `{}` has been assigned.", role.name))
.await?;
}
}
};
@ -144,27 +156,29 @@ fn toggle(ctx: &mut Context, m: &Message, mut args: Args) -> CommandResult {
#[num_args(2)]
#[required_permissions(MANAGE_ROLES)]
#[only_in(guilds)]
fn add(ctx: &mut Context, m: &Message, mut args: Args) -> CommandResult {
async fn add(ctx: &Context, m: &Message, mut args: Args) -> CommandResult {
let role = args.single_quoted::<String>()?;
let data = ctx.data.read().await;
let description = args.single::<String>()?;
let guild_id = m.guild_id.unwrap();
let roles = guild_id.to_partial_guild(&ctx)?.roles;
let roles = guild_id.to_partial_guild(&ctx).await?.roles;
let role = role_from_string(&role, &roles);
match role {
None => {
m.reply(&ctx, "No such role exists")?;
m.reply(&ctx, "No such role exists").await?;
}
Some(role)
if DB::open(&*ctx.data.read())
if DB::open(&*data)
.borrow()?
.get(&guild_id)
.map(|g| g.contains_key(&role.id))
.unwrap_or(false) =>
{
m.reply(&ctx, "This role already exists in the database.")?;
m.reply(&ctx, "This role already exists in the database.")
.await?;
}
Some(role) => {
DB::open(&*ctx.data.read())
DB::open(&*data)
.borrow_mut()?
.entry(guild_id)
.or_default()
@ -175,7 +189,7 @@ fn add(ctx: &mut Context, m: &Message, mut args: Args) -> CommandResult {
description,
},
);
m.react(&ctx, "👌🏼")?;
m.react(&ctx, '👌').await?;
}
};
Ok(())
@ -188,31 +202,33 @@ fn add(ctx: &mut Context, m: &Message, mut args: Args) -> CommandResult {
#[num_args(1)]
#[required_permissions(MANAGE_ROLES)]
#[only_in(guilds)]
fn remove(ctx: &mut Context, m: &Message, mut args: Args) -> CommandResult {
async fn remove(ctx: &Context, m: &Message, mut args: Args) -> CommandResult {
let role = args.single_quoted::<String>()?;
let data = ctx.data.read().await;
let guild_id = m.guild_id.unwrap();
let roles = guild_id.to_partial_guild(&ctx)?.roles;
let roles = guild_id.to_partial_guild(&ctx).await?.roles;
let role = role_from_string(&role, &roles);
match role {
None => {
m.reply(&ctx, "No such role exists")?;
m.reply(&ctx, "No such role exists").await?;
}
Some(role)
if !DB::open(&*ctx.data.read())
if !DB::open(&*data)
.borrow()?
.get(&guild_id)
.map(|g| g.contains_key(&role.id))
.unwrap_or(false) =>
{
m.reply(&ctx, "This role does not exist in the assignable list.")?;
m.reply(&ctx, "This role does not exist in the assignable list.")
.await?;
}
Some(role) => {
DB::open(&*ctx.data.read())
DB::open(&*data)
.borrow_mut()?
.entry(guild_id)
.or_default()
.remove(&role.id);
m.react(&ctx, "👌🏼")?;
m.react(&ctx, '👌').await?;
}
};
Ok(())

214
youmubot-core/src/community/votes.rs
View file

@ -1,14 +1,18 @@
use serenity::framework::standard::CommandError as Error;
use serenity::{
collector::ReactionAction,
framework::standard::{macros::command, Args, CommandResult},
model::{
channel::{Message, Reaction, ReactionType},
channel::{Message, ReactionType},
id::UserId,
},
utils::MessageBuilder,
};
use std::collections::{HashMap as Map, HashSet as Set};
use std::time::Duration;
use std::{
collections::{HashMap as Map, HashSet as Set},
convert::TryFrom,
};
use youmubot_prelude::{Duration as ParseDuration, *};
#[command]
@ -19,13 +23,13 @@ use youmubot_prelude::{Duration as ParseDuration, *};
#[only_in(guilds)]
#[min_args(2)]
#[owner_privilege]
pub fn vote(ctx: &mut Context, msg: &Message, mut args: Args) -> CommandResult {
pub async fn vote(ctx: &Context, msg: &Message, mut args: Args) -> CommandResult {
// Parse stuff first
let args = args.quoted();
let _duration = args.single::<ParseDuration>()?;
let duration = &_duration.0;
if *duration < Duration::from_secs(2 * 60) || *duration > Duration::from_secs(60 * 60 * 24) {
msg.reply(ctx, format!("😒 Invalid duration ({}). The voting time should be between **2 minutes** and **1 day**.", _duration))?;
if *duration < Duration::from_secs(2) || *duration > Duration::from_secs(60 * 60 * 24) {
msg.reply(ctx, format!("😒 Invalid duration ({}). The voting time should be between **2 minutes** and **1 day**.", _duration)).await?;
return Ok(());
}
let question = args.single::<String>()?;
@ -41,7 +45,8 @@ pub fn vote(ctx: &mut Context, msg: &Message, mut args: Args) -> CommandResult {
msg.reply(
ctx,
"😒 Can't have a nice voting session if you only have one choice.",
)?;
)
.await?;
return Ok(());
}
if choices.len() > MAX_CHOICES {
@ -52,7 +57,8 @@ pub fn vote(ctx: &mut Context, msg: &Message, mut args: Args) -> CommandResult {
"😵 Too many choices... We only support {} choices at the moment!",
MAX_CHOICES
),
)?;
)
.await?;
return Ok(());
}
@ -89,123 +95,116 @@ pub fn vote(ctx: &mut Context, msg: &Message, mut args: Args) -> CommandResult {
.description(MessageBuilder::new().push_bold_line_safe(&question).push("\nThis question was asked by ").push(author.mention()))
.fields(fields.into_iter())
})
})?;
msg.delete(&ctx)?;
}).await?;
msg.delete(&ctx).await?;
drop(msg);
// React on all the choices
choices
.iter()
.try_for_each(|(emote, _)| panel.react(&ctx, &emote[..]))?;
.map(|(emote, _)| {
panel
.react(&ctx, ReactionType::try_from(&emote[..]).unwrap())
.map_ok(|_| ())
})
.collect::<stream::FuturesUnordered<_>>()
.try_collect::<()>()
.await?;
// A handler for votes.
struct VoteHandler {
pub ctx: Context,
pub msg: Message,
pub user_reactions: Map<String, Set<UserId>>,
let user_reactions: Map<String, Set<UserId>> = choices
.iter()
.map(|(emote, _)| (emote.clone(), Set::new()))
.collect();
pub panel: Message,
}
impl VoteHandler {
fn new(ctx: Context, msg: Message, panel: Message, choices: &[(String, String)]) -> Self {
VoteHandler {
ctx,
msg,
user_reactions: choices
.iter()
.map(|(emote, _)| (emote.clone(), Set::new()))
.collect(),
panel,
}
}
}
impl ReactionHandler for VoteHandler {
fn handle_reaction(&mut self, reaction: &Reaction, is_add: bool) -> CommandResult {
if reaction.message_id != self.panel.id {
return Ok(());
}
if reaction.user(&self.ctx)?.bot {
return Ok(());
}
// Collect reactions...
let user_reactions = panel
.await_reactions(&ctx)
.removed(true)
.timeout(*duration)
.await
.fold(user_reactions, |mut set, reaction| async move {
let (reaction, is_add) = match &*reaction {
ReactionAction::Added(r) => (r, true),
ReactionAction::Removed(r) => (r, false),
};
let users = if let ReactionType::Unicode(ref s) = reaction.emoji {
if let Some(users) = self.user_reactions.get_mut(s.as_str()) {
if let Some(users) = set.get_mut(s.as_str()) {
users
} else {
return Ok(());
return set;
}
} else {
return Ok(());
return set;
};
let user_id = match reaction.user_id {
Some(v) => v,
None => return set,
};
if is_add {
users.insert(reaction.user_id);
users.insert(user_id);
} else {
users.remove(&reaction.user_id);
users.remove(&user_id);
}
Ok(())
}
set
})
.await;
// Handle choices
let choice_map = choices.into_iter().collect::<Map<_, _>>();
let mut result: Vec<(String, Vec<UserId>)> = user_reactions
.into_iter()
.filter(|(_, users)| !users.is_empty())
.map(|(emote, users)| (emote, users.into_iter().collect()))
.collect();
result.sort_unstable_by(|(_, v), (_, w)| w.len().cmp(&v.len()));
if result.len() == 0 {
msg.reply(
&ctx,
MessageBuilder::new()
.push("no one answer your question ")
.push_bold_safe(&question)
.push(", sorry 😭")
.build(),
)
.await?;
return Ok(());
}
ctx.data
.get_cloned::<ReactionWatcher>()
.handle_reactions_timed(
VoteHandler::new(ctx.clone(), msg.clone(), panel, &choices),
*duration,
move |vh| {
let (ctx, msg, user_reactions, panel) =
(vh.ctx, vh.msg, vh.user_reactions, vh.panel);
let choice_map = choices.into_iter().collect::<Map<_, _>>();
let result: Vec<(String, Vec<UserId>)> = user_reactions
.into_iter()
.filter(|(_, users)| !users.is_empty())
.map(|(emote, users)| (emote, users.into_iter().collect()))
.collect();
if result.len() == 0 {
msg.reply(
&ctx,
MessageBuilder::new()
.push("no one answer your question ")
.push_bold_safe(&question)
.push(", sorry 😭")
.build(),
)
.ok();
} else {
channel
.send_message(&ctx, |c| {
c.content({
let mut content = MessageBuilder::new();
content
.push("@here, ")
.push(author.mention())
.push(" previously asked ")
.push_bold_safe(&question)
.push(", and here are the results!");
result.into_iter().for_each(|(emote, votes)| {
content
.push("\n - ")
.push_bold(format!("{}", votes.len()))
.push(" voted for ")
.push(&emote)
.push(" ")
.push_bold_safe(choice_map.get(&emote).unwrap())
.push(": ")
.push(
votes
.into_iter()
.map(|v| v.mention())
.collect::<Vec<_>>()
.join(", "),
);
});
content.build()
})
})
.ok();
}
panel.delete(&ctx).ok();
},
);
channel
.send_message(&ctx, |c| {
c.content({
let mut content = MessageBuilder::new();
content
.push("@here, ")
.push(author.mention())
.push(" previously asked ")
.push_bold_safe(&question)
.push(", and here are the results!");
result.into_iter().for_each(|(emote, votes)| {
content
.push("\n - ")
.push_bold(format!("{}", votes.len()))
.push(" voted for ")
.push(&emote)
.push(" ")
.push_bold_safe(choice_map.get(&emote).unwrap())
.push(": ")
.push(
votes
.into_iter()
.map(|v| v.mention())
.collect::<Vec<_>>()
.join(", "),
);
});
content.build()
})
})
.await?;
panel.delete(&ctx).await?;
Ok(())
// unimplemented!();
@ -239,3 +238,4 @@ const REACTIONS: [&'static str; 90] = [
// Assertions
static_assertions::const_assert!(MAX_CHOICES <= REACTIONS.len());
static_assertions::const_assert!(MAX_CHOICES <= REACTIONS.len());

27
youmubot-core/src/db.rs
View file

@ -14,30 +14,25 @@ pub type Roles = DB<GuildMap<HashMap<RoleId, Role>>>;
/// For the admin commands:
/// - Each server might have a `soft ban` role implemented.
/// - We allow periodical `soft ban` applications.
#[derive(Serialize, Deserialize, Debug, Clone)]
pub enum ServerSoftBans {
Implemented(ImplementedSoftBans),
Unimplemented,
}
impl ServerSoftBans {
// Create a new, implemented role.
pub fn new_implemented(role: RoleId) -> ServerSoftBans {
ServerSoftBans::Implemented(ImplementedSoftBans {
role,
periodical_bans: HashMap::new(),
})
}
}
#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct ImplementedSoftBans {
pub struct ServerSoftBans {
/// The soft-ban role.
pub role: RoleId,
/// List of all to-unban people.
pub periodical_bans: HashMap<UserId, DateTime<Utc>>,
}
impl ServerSoftBans {
// Create a new, implemented role.
pub fn new(role: RoleId) -> Self {
Self {
role,
periodical_bans: HashMap::new(),
}
}
}
/// Role represents an assignable role.
#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct Role {

37
youmubot-core/src/fun/images.rs
View file

@ -15,24 +15,24 @@ use youmubot_prelude::*;
#[description = "🖼️ Find an image with a given tag on Danbooru[nsfw]!"]
#[min_args(1)]
#[bucket("images")]
pub fn nsfw(ctx: &mut Context, msg: &Message, args: Args) -> CommandResult {
message_command(ctx, msg, args, Rating::Explicit)
pub async fn nsfw(ctx: &Context, msg: &Message, args: Args) -> CommandResult {
message_command(ctx, msg, args, Rating::Explicit).await
}
#[command]
#[description = "🖼️ Find an image with a given tag on Danbooru[safe]!"]
#[min_args(1)]
#[bucket("images")]
pub fn image(ctx: &mut Context, msg: &Message, args: Args) -> CommandResult {
message_command(ctx, msg, args, Rating::Safe)
pub async fn image(ctx: &Context, msg: &Message, args: Args) -> CommandResult {
message_command(ctx, msg, args, Rating::Safe).await
}
#[check]
#[name = "nsfw"]
fn nsfw_check(ctx: &mut Context, msg: &Message, _: &mut Args, _: &CommandOptions) -> CheckResult {
let channel = msg.channel_id.to_channel(&ctx).unwrap();
async fn nsfw_check(ctx: &Context, msg: &Message, _: &mut Args, _: &CommandOptions) -> CheckResult {
let channel = msg.channel_id.to_channel(&ctx).await.unwrap();
if !(match channel {
Channel::Guild(guild_channel) => guild_channel.read().nsfw,
Channel::Guild(guild_channel) => guild_channel.nsfw,
_ => true,
}) {
CheckResult::Failure(Reason::User("😣 YOU FREAKING PERVERT!!!".to_owned()))
@ -41,22 +41,31 @@ fn nsfw_check(ctx: &mut Context, msg: &Message, _: &mut Args, _: &CommandOptions
}
}
fn message_command(ctx: &mut Context, msg: &Message, args: Args, rating: Rating) -> CommandResult {
async fn message_command(
ctx: &Context,
msg: &Message,
args: Args,
rating: Rating,
) -> CommandResult {
let tags = args.remains().unwrap_or("touhou");
let http = ctx.data.get_cloned::<HTTPClient>();
let image = get_image(&http, rating, tags)?;
let image = get_image(
ctx.data.read().await.get::<HTTPClient>().unwrap(),
rating,
tags,
)
.await?;
match image {
None => msg.reply(&ctx, "🖼️ No image found...\n💡 Tip: In danbooru, character names follow Japanese standards (last name before first name), so **Hakurei Reimu** might give you an image while **Reimu Hakurei** won't."),
None => msg.reply(&ctx, "🖼️ No image found...\n💡 Tip: In danbooru, character names follow Japanese standards (last name before first name), so **Hakurei Reimu** might give you an image while **Reimu Hakurei** won't.").await,
Some(url) => msg.reply(
&ctx,
format!("🖼️ Here's the image you requested!\n\n{}", url),
),
).await,
}?;
Ok(())
}
// Gets an image URL.
fn get_image(
async fn get_image(
client: &<HTTPClient as TypeMapKey>::Value,
rating: Rating,
tags: &str,
@ -72,7 +81,7 @@ fn get_image(
.query(&[("limit", "1"), ("random", "true")])
.build()?;
println!("{:?}", req.url());
let response: Vec<PostResponse> = client.execute(req)?.json()?;
let response: Vec<PostResponse> = client.execute(req).await?.json().await?;
Ok(response
.into_iter()
.next()

100
youmubot-core/src/fun/mod.rs
View file

@ -28,7 +28,7 @@ struct Fun;
#[max_args(2)]
#[usage = "[max-dice-faces = 6] / [message]"]
#[example = "100 / What's my score?"]
fn roll(ctx: &mut Context, msg: &Message, mut args: Args) -> CommandResult {
async fn roll(ctx: &Context, msg: &Message, mut args: Args) -> CommandResult {
let dice = if args.is_empty() {
6
} else {
@ -36,7 +36,8 @@ fn roll(ctx: &mut Context, msg: &Message, mut args: Args) -> CommandResult {
};
if dice == 0 {
msg.reply(&ctx, "Give me a dice with 0 faces, what do you expect 😒")?;
msg.reply(&ctx, "Give me a dice with 0 faces, what do you expect 😒")
.await?;
return Ok(());
}
@ -47,24 +48,30 @@ fn roll(ctx: &mut Context, msg: &Message, mut args: Args) -> CommandResult {
};
match args.single_quoted::<String>() {
Ok(s) => msg.reply(
&ctx,
MessageBuilder::new()
.push("you asked ")
.push_bold_safe(s)
.push(format!(
", so I rolled a 🎲 of **{}** faces, and got **{}**!",
Ok(s) => {
msg.reply(
&ctx,
MessageBuilder::new()
.push("you asked ")
.push_bold_safe(s)
.push(format!(
", so I rolled a 🎲 of **{}** faces, and got **{}**!",
dice, result
))
.build(),
)
.await
}
Err(_) if args.is_empty() => {
msg.reply(
&ctx,
format!(
"I rolled a 🎲 of **{}** faces, and got **{}**!",
dice, result
))
.build(),
),
Err(_) if args.is_empty() => msg.reply(
&ctx,
format!(
"I rolled a 🎲 of **{}** faces, and got **{}**!",
dice, result
),
),
),
)
.await
}
Err(e) => return Err(e.into()),
}?;
@ -77,7 +84,7 @@ You may prefix the first choice with `?` to make it a question!
If no choices are given, Youmu defaults to `Yes!` and `No!`"#]
#[usage = "[?question]/[choice #1]/[choice #2]/..."]
#[example = "?What for dinner/Pizza/Hamburger"]
fn pick(ctx: &mut Context, msg: &Message, mut args: Args) -> CommandResult {
async fn pick(ctx: &Context, msg: &Message, mut args: Args) -> CommandResult {
let (question, choices) = {
// Get a list of options.
let mut choices = args
@ -114,24 +121,30 @@ fn pick(ctx: &mut Context, msg: &Message, mut args: Args) -> CommandResult {
};
match question {
None => msg.reply(
&ctx,
MessageBuilder::new()
.push("Youmu picks 👉")
.push_bold_safe(choice)
.push("👈!")
.build(),
),
Some(s) => msg.reply(
&ctx,
MessageBuilder::new()
.push("you asked ")
.push_bold_safe(s)
.push(", and Youmu picks 👉")
.push_bold_safe(choice)
.push("👈!")
.build(),
),
None => {
msg.reply(
&ctx,
MessageBuilder::new()
.push("Youmu picks 👉")
.push_bold_safe(choice)
.push("👈!")
.build(),
)
.await
}
Some(s) => {
msg.reply(
&ctx,
MessageBuilder::new()
.push("you asked ")
.push_bold_safe(s)
.push(", and Youmu picks 👉")
.push_bold_safe(choice)
.push("👈!")
.build(),
)
.await
}
}?;
Ok(())
@ -142,7 +155,7 @@ fn pick(ctx: &mut Context, msg: &Message, mut args: Args) -> CommandResult {
#[usage = "[user_mention = yourself]"]
#[example = "@user#1234"]
#[max_args(1)]
fn name(ctx: &mut Context, msg: &Message, mut args: Args) -> CommandResult {
async fn name(ctx: &Context, msg: &Message, mut args: Args) -> CommandResult {
let user_id = if args.is_empty() {
msg.author.id
} else {
@ -153,15 +166,15 @@ fn name(ctx: &mut Context, msg: &Message, mut args: Args) -> CommandResult {
"your".to_owned()
} else {
MessageBuilder::new()
.push_bold_safe(user_id.to_user(&ctx)?.tag())
.push_bold_safe(user_id.to_user(&ctx).await?.tag())
.push("'s")
.build()
};
// Rule out a couple of cases
if user_id == ctx.http.get_current_application_info()?.id {
if user_id == ctx.http.get_current_application_info().await?.id {
// This is my own user_id
msg.reply(&ctx, "😠 My name is **Youmu Konpaku**!")?;
msg.reply(&ctx, "😠 My name is **Youmu Konpaku**!").await?;
return Ok(());
}
@ -173,6 +186,7 @@ fn name(ctx: &mut Context, msg: &Message, mut args: Args) -> CommandResult {
"{} Japanese🇯🇵 name is **{} {}**!",
user_mention, first_name, last_name
),
)?;
)
.await?;
Ok(())
}

14
youmubot-core/src/lib.rs
View file

@ -20,26 +20,30 @@ pub use fun::FUN_GROUP;
pub fn setup(
path: &std::path::Path,
client: &serenity::client::Client,
data: &mut youmubot_prelude::ShareMap,
data: &mut TypeMap,
) -> serenity::framework::standard::CommandResult {
db::SoftBans::insert_into(&mut *data, &path.join("soft_bans.yaml"))?;
db::Roles::insert_into(&mut *data, &path.join("roles.yaml"))?;
// Create handler threads
std::thread::spawn(admin::watch_soft_bans(client));
tokio::spawn(admin::watch_soft_bans(
client.cache_and_http.clone(),
client.data.clone(),
));
Ok(())
}
// A help command
#[help]
pub fn help(
context: &mut Context,
pub async fn help(
context: &Context,
msg: &Message,
args: Args,
help_options: &'static HelpOptions,
groups: &[&'static CommandGroup],
owners: HashSet<UserId>,
) -> CommandResult {
help_commands::with_embeds(context, msg, args, help_options, groups, owners)
help_commands::with_embeds(context, msg, args, help_options, groups, owners).await;
Ok(())
}

11
youmubot-db/Cargo.toml
View file

@ -7,18 +7,11 @@ edition = "2018"
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
[dependencies]
serenity = "0.8"
serenity = "0.9.0-rc.0"
dotenv = "0.15"
serde = { version = "1.0", features = ["derive"] }
chrono = "0.4.9"
# rand = "0.7.2"
# static_assertions = "1.1.0"
# reqwest = "0.10.1"
# regex = "1"
# lazy_static = "1"
# youmubot-osu = { path = "../youmubot-osu" }
rayon = "1.1"
[dependencies.rustbreak]
version = "2.0.0-rc3"
version = "2.0.0"
features = ["yaml_enc"]

61
youmubot-db/src/lib.rs
View file

@ -1,15 +1,22 @@
use rustbreak::{deser::Yaml as Ron, FileDatabase};
use rustbreak::{deser::Yaml, FileDatabase, RustbreakError as DBError};
use serde::{de::DeserializeOwned, Deserialize, Serialize};
use serenity::{framework::standard::CommandError as Error, model::id::GuildId, prelude::*};
use std::collections::HashMap;
use std::{cell::Cell, path::Path, sync::Arc};
use serenity::{
model::id::GuildId,
prelude::{TypeMap, TypeMapKey},
};
use std::{collections::HashMap, path::Path};
/// GuildMap defines the guild-map type.
/// It is basically a HashMap from a GuildId to a data structure.
pub type GuildMap<V> = HashMap<GuildId, V>;
/// The generic DB type we will be using.
pub struct DB<T>(std::marker::PhantomData<T>);
impl<T: std::any::Any> serenity::prelude::TypeMapKey for DB<T> {
type Value = Arc<FileDatabase<T, Ron>>;
/// A short type abbreviation for a FileDatabase.
type Database<T> = FileDatabase<T, Yaml>;
impl<T: std::any::Any + Send + Sync> TypeMapKey for DB<T> {
type Value = Database<T>;
}
impl<T: std::any::Any + Default + Send + Sync + Clone + Serialize + std::fmt::Debug> DB<T>
@ -17,66 +24,62 @@ where
for<'de> T: Deserialize<'de>,
{
/// Insert into a ShareMap.
pub fn insert_into(data: &mut ShareMap, path: impl AsRef<Path>) -> Result<(), Error> {
let db = FileDatabase::<T, Ron>::from_path(path, T::default())?;
db.load().or_else(|e| {
dbg!(e);
db.save()
})?;
data.insert::<DB<T>>(Arc::new(db));
pub fn insert_into(data: &mut TypeMap, path: impl AsRef<Path>) -> Result<(), DBError> {
let db = Database::<T>::load_from_path_or_default(path)?;
data.insert::<DB<T>>(db);
Ok(())
}
/// Open a previously inserted DB.
pub fn open(data: &ShareMap) -> DBWriteGuard<T> {
data.get::<Self>().expect("DB initialized").clone().into()
pub fn open(data: &TypeMap) -> DBWriteGuard<T> {
data.get::<Self>().expect("DB initialized").into()
}
}
/// The write guard for our FileDatabase.
/// It wraps the FileDatabase in a write-on-drop lock.
#[derive(Debug)]
pub struct DBWriteGuard<T>
pub struct DBWriteGuard<'a, T>
where
T: Send + Sync + Clone + std::fmt::Debug + Serialize + DeserializeOwned,
{
db: Arc<FileDatabase<T, Ron>>,
needs_save: Cell<bool>,
db: &'a Database<T>,
needs_save: bool,
}
impl<T> From<Arc<FileDatabase<T, Ron>>> for DBWriteGuard<T>
impl<'a, T> From<&'a Database<T>> for DBWriteGuard<'a, T>
where
T: Send + Sync + Clone + std::fmt::Debug + Serialize + DeserializeOwned,
{
fn from(v: Arc<FileDatabase<T, Ron>>) -> Self {
fn from(v: &'a Database<T>) -> Self {
DBWriteGuard {
db: v,
needs_save: Cell::new(false),
needs_save: false,
}
}
}
impl<T> DBWriteGuard<T>
impl<'a, T> DBWriteGuard<'a, T>
where
T: Send + Sync + Clone + std::fmt::Debug + Serialize + DeserializeOwned,
{
/// Borrows the FileDatabase.
pub fn borrow(&self) -> Result<std::sync::RwLockReadGuard<T>, rustbreak::RustbreakError> {
(*self).db.borrow_data()
pub fn borrow(&self) -> Result<std::sync::RwLockReadGuard<T>, DBError> {
self.db.borrow_data()
}
/// Borrows the FileDatabase for writing.
pub fn borrow_mut(&self) -> Result<std::sync::RwLockWriteGuard<T>, rustbreak::RustbreakError> {
self.needs_save.set(true);
(*self).db.borrow_data_mut()
pub fn borrow_mut(&mut self) -> Result<std::sync::RwLockWriteGuard<T>, DBError> {
self.needs_save = true;
self.db.borrow_data_mut()
}
}
impl<T> Drop for DBWriteGuard<T>
impl<'a, T> Drop for DBWriteGuard<'a, T>
where
T: Send + Sync + Clone + std::fmt::Debug + Serialize + DeserializeOwned,
{
fn drop(&mut self) {
if self.needs_save.get() {
if self.needs_save {
if let Err(e) = self.db.save() {
dbg!(e);
}

3
youmubot-osu/Cargo.toml
View file

@ -7,12 +7,11 @@ edition = "2018"
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
[dependencies]
serenity = "0.8"
serenity = "0.9.0-rc.0"
chrono = "0.4.10"
reqwest = "0.10.1"
serde = { version = "1.0", features = ["derive"] }
bitflags = "1"
rayon = "1.1"
lazy_static = "1"
regex = "1"
oppai-rs = "0.2.0"

200
youmubot-osu/src/discord/announcer.rs
View file

@ -9,101 +9,159 @@ use crate::{
Client as Osu,
};
use announcer::MemberToChannels;
use rayon::prelude::*;
use serenity::{
framework::standard::{CommandError as Error, CommandResult},
http::CacheHttp,
model::id::{ChannelId, UserId},
CacheAndHttp,
};
use std::sync::Arc;
use std::{collections::HashMap, sync::Arc};
use youmubot_prelude::*;
/// osu! announcer's unique announcer key.
pub const ANNOUNCER_KEY: &'static str = "osu";
/// Announce osu! top scores.
pub fn updates(c: Arc<CacheAndHttp>, d: AppData, channels: MemberToChannels) -> CommandResult {
let osu = d.get_cloned::<OsuClient>();
let cache = d.get_cloned::<BeatmapMetaCache>();
let oppai = d.get_cloned::<BeatmapCache>();
// For each user...
let mut data = OsuSavedUsers::open(&*d.read()).borrow()?.clone();
for (user_id, osu_user) in data.iter_mut() {
let channels = channels.channels_of(c.clone(), *user_id);
if channels.is_empty() {
continue; // We don't wanna update an user without any active server
}
osu_user.pp = match (&[Mode::Std, Mode::Taiko, Mode::Catch, Mode::Mania])
.par_iter()
.map(|m| {
handle_user_mode(
c.clone(),
&osu,
&cache,
&oppai,
&osu_user,
*user_id,
&channels[..],
*m,
d.clone(),
)
/// The announcer struct implementing youmubot_prelude::Announcer
pub struct Announcer;
#[async_trait]
impl youmubot_prelude::Announcer for Announcer {
async fn updates(
&mut self,
c: Arc<CacheAndHttp>,
d: AppData,
channels: MemberToChannels,
) -> Result<()> {
// For each user...
let data = OsuSavedUsers::open(&*d.read().await).borrow()?.clone();
let data = data
.into_iter()
.map(|(user_id, osu_user)| {
let d = d.clone();
let channels = &channels;
let c = c.clone();
async move {
let channels = channels.channels_of(c.clone(), user_id).await;
if channels.is_empty() {
return (user_id, osu_user); // We don't wanna update an user without any active server
}
let pp = match (&[Mode::Std, Mode::Taiko, Mode::Catch, Mode::Mania])
.into_iter()
.map(|m| {
handle_user_mode(
c.clone(),
&osu_user,
user_id,
channels.clone(),
*m,
d.clone(),
)
})
.collect::<stream::FuturesOrdered<_>>()
.try_collect::<Vec<_>>()
.await
{
Ok(v) => v,
Err(e) => {
eprintln!("osu: Cannot update {}: {}", osu_user.id, e);
return (user_id, osu_user);
}
};
let last_update = chrono::Utc::now();
(
user_id,
OsuUser {
pp,
last_update,
..osu_user
},
)
}
})
.collect::<Result<_, _>>()
{
Ok(v) => v,
Err(e) => {
eprintln!("osu: Cannot update {}: {}", osu_user.id, e.0);
continue;
}
};
osu_user.last_update = chrono::Utc::now();
.collect::<stream::FuturesUnordered<_>>()
.collect::<HashMap<_, _>>()
.await;
// Update users
*OsuSavedUsers::open(&*d.read().await).borrow_mut()? = data;
Ok(())
}
// Update users
*OsuSavedUsers::open(&*d.read()).borrow_mut()? = data;
Ok(())
}
/// Handles an user/mode scan, announces all possible new scores, return the new pp value.
fn handle_user_mode(
async fn handle_user_mode(
c: Arc<CacheAndHttp>,
osu: &Osu,
cache: &BeatmapMetaCache,
oppai: &BeatmapCache,
osu_user: &OsuUser,
user_id: UserId,
channels: &[ChannelId],
channels: Vec<ChannelId>,
mode: Mode,
d: AppData,
) -> Result<Option<f64>, Error> {
let scores = scan_user(osu, osu_user, mode)?;
let user = osu
.user(UserID::ID(osu_user.id), |f| f.mode(mode))?
.ok_or(Error::from("user not found"))?;
scores
.into_par_iter()
.map(|(rank, score)| -> Result<_, Error> {
let beatmap = cache.get_beatmap_default(score.beatmap_id)?;
let content = oppai.get_beatmap(beatmap.beatmap_id)?;
Ok((rank, score, BeatmapWithMode(beatmap, mode), content))
})
.filter_map(|v| v.ok())
.for_each(|(rank, score, beatmap, content)| {
for channel in (&channels).iter() {
if let Err(e) = channel.send_message(c.http(), |c| {
c.content(format!("New top record from {}!", user_id.mention()))
.embed(|e| score_embed(&score, &beatmap, &content, &user, Some(rank), e))
}) {
dbg!(e);
let (scores, user) = {
let data = d.read().await;
let osu = data.get::<OsuClient>().unwrap();
let scores = scan_user(osu, osu_user, mode).await?;
let user = osu
.user(UserID::ID(osu_user.id), |f| f.mode(mode))
.await?
.ok_or(Error::msg("user not found"))?;
(scores, user)
};
let pp = user.pp;
spawn_future(async move {
scores
.into_iter()
.map(|(rank, score)| {
let d = d.clone();
async move {
let data = d.read().await;
let cache = data.get::<BeatmapMetaCache>().unwrap();
let oppai = data.get::<BeatmapCache>().unwrap();
let beatmap = cache.get_beatmap_default(score.beatmap_id).await?;
let content = oppai.get_beatmap(beatmap.beatmap_id).await?;
let r: Result<_> = Ok((rank, score, BeatmapWithMode(beatmap, mode), content));
r
}
save_beatmap(&*d.read(), *channel, &beatmap).ok();
}
});
Ok(user.pp)
})
.collect::<stream::FuturesOrdered<_>>()
.filter_map(|v| future::ready(v.ok()))
.for_each(move |(rank, score, beatmap, content)| {
let channels = channels.clone();
let d = d.clone();
let c = c.clone();
let user = user.clone();
async move {
let data = d.read().await;
for channel in (&channels).iter() {
if let Err(e) = channel
.send_message(c.http(), |c| {
c.content(format!("New top record from {}!", user_id.mention()))
.embed(|e| {
score_embed(
&score,
&beatmap,
&content,
&user,
Some(rank),
e,
)
})
})
.await
{
dbg!(e);
}
save_beatmap(&*data, *channel, &beatmap).ok();
}
}
})
.await;
});
Ok(pp)
}
fn scan_user(osu: &Osu, u: &OsuUser, mode: Mode) -> Result<Vec<(u8, Score)>, Error> {
let scores = osu.user_best(UserID::ID(u.id), |f| f.mode(mode).limit(25))?;
async fn scan_user(osu: &Osu, u: &OsuUser, mode: Mode) -> Result<Vec<(u8, Score)>, Error> {
let scores = osu
.user_best(UserID::ID(u.id), |f| f.mode(mode).limit(25))
.await?;
let scores = scores
.into_iter()
.enumerate()

58
youmubot-osu/src/discord/beatmap_cache.rs
View file

@ -3,16 +3,14 @@ use crate::{
Client,
};
use dashmap::DashMap;
use serenity::framework::standard::CommandError;
use std::sync::Arc;
use youmubot_prelude::TypeMapKey;
use youmubot_prelude::*;
/// BeatmapMetaCache intercepts beatmap-by-id requests and caches them for later recalling.
/// Does not cache non-Ranked beatmaps.
#[derive(Clone, Debug)]
pub struct BeatmapMetaCache {
client: Client,
cache: Arc<DashMap<(u64, Mode), Beatmap>>,
client: Arc<Client>,
cache: DashMap<(u64, Mode), Beatmap>,
}
impl TypeMapKey for BeatmapMetaCache {
@ -21,13 +19,13 @@ impl TypeMapKey for BeatmapMetaCache {
impl BeatmapMetaCache {
/// Create a new beatmap cache.
pub fn new(client: Client) -> Self {
pub fn new(client: Arc<Client>) -> Self {
BeatmapMetaCache {
client,
cache: Arc::new(DashMap::new()),
cache: DashMap::new(),
}
}
fn insert_if_possible(&self, id: u64, mode: Option<Mode>) -> Result<Beatmap, CommandError> {
async fn insert_if_possible(&self, id: u64, mode: Option<Mode>) -> Result<Beatmap> {
let beatmap = self
.client
.beatmaps(crate::BeatmapRequestKind::Beatmap(id), |f| {
@ -36,35 +34,37 @@ impl BeatmapMetaCache {
}
f
})
.and_then(|v| {
v.into_iter()
.next()
.ok_or(CommandError::from("beatmap not found"))
})?;
.await
.and_then(|v| v.into_iter().next().ok_or(Error::msg("beatmap not found")))?;
if let ApprovalStatus::Ranked(_) = beatmap.approval {
self.cache.insert((id, beatmap.mode), beatmap.clone());
};
Ok(beatmap)
}
/// Get the given beatmap
pub fn get_beatmap(&self, id: u64, mode: Mode) -> Result<Beatmap, CommandError> {
self.cache
.get(&(id, mode))
.map(|b| Ok(b.clone()))
.unwrap_or_else(|| self.insert_if_possible(id, Some(mode)))
pub async fn get_beatmap(&self, id: u64, mode: Mode) -> Result<Beatmap> {
match self.cache.get(&(id, mode)).map(|v| v.clone()) {
Some(v) => Ok(v),
None => self.insert_if_possible(id, Some(mode)).await,
}
}
/// Get a beatmap without a mode...
pub fn get_beatmap_default(&self, id: u64) -> Result<Beatmap, CommandError> {
(&[Mode::Std, Mode::Taiko, Mode::Catch, Mode::Mania])
.iter()
.filter_map(|&mode| {
self.cache
.get(&(id, mode))
.filter(|b| b.mode == mode)
.map(|b| Ok(b.clone()))
})
.next()
.unwrap_or_else(|| self.insert_if_possible(id, None))
pub async fn get_beatmap_default(&self, id: u64) -> Result<Beatmap> {
Ok(
match (&[Mode::Std, Mode::Taiko, Mode::Catch, Mode::Mania])
.iter()
.filter_map(|&mode| {
self.cache
.get(&(id, mode))
.filter(|b| b.mode == mode)
.map(|b| b.clone())
})
.next()
{
Some(v) => v,
None => self.insert_if_possible(id, None).await?,
},
)
}
}

22
youmubot-osu/src/discord/cache.rs
View file

@ -1,30 +1,26 @@
use super::db::OsuLastBeatmap;
use super::BeatmapWithMode;
use serenity::{
framework::standard::{CommandError as Error, CommandResult},
model::id::ChannelId,
prelude::*,
};
use serenity::model::id::ChannelId;
use youmubot_prelude::*;
/// Save the beatmap into the server data storage.
pub(crate) fn save_beatmap(
data: &ShareMap,
data: &TypeMap,
channel_id: ChannelId,
bm: &BeatmapWithMode,
) -> CommandResult {
let db = OsuLastBeatmap::open(data);
let mut db = db.borrow_mut()?;
db.insert(channel_id, (bm.0.clone(), bm.mode()));
) -> Result<()> {
OsuLastBeatmap::open(data)
.borrow_mut()?
.insert(channel_id, (bm.0.clone(), bm.mode()));
Ok(())
}
/// Get the last beatmap requested from this channel.
pub(crate) fn get_beatmap(
data: &ShareMap,
data: &TypeMap,
channel_id: ChannelId,
) -> Result<Option<BeatmapWithMode>, Error> {
) -> Result<Option<BeatmapWithMode>> {
let db = OsuLastBeatmap::open(data);
let db = db.borrow()?;

363
youmubot-osu/src/discord/hook.rs
View file

@ -7,12 +7,7 @@ use crate::{
};
use lazy_static::lazy_static;
use regex::Regex;
use serenity::{
builder::CreateMessage,
framework::standard::{CommandError as Error, CommandResult},
model::channel::Message,
utils::MessageBuilder,
};
use serenity::{builder::CreateMessage, model::channel::Message, utils::MessageBuilder};
use std::str::FromStr;
use youmubot_prelude::*;
@ -26,47 +21,58 @@ lazy_static! {
r"(?:https?://)?osu\.ppy\.sh/beatmapsets/(?P<set_id>\d+)/?(?:\#(?P<mode>osu|taiko|fruits|mania)(?:/(?P<beatmap_id>\d+)|/?))?(?:\+(?P<mods>[A-Z]+))?"
).unwrap();
static ref SHORT_LINK_REGEX: Regex = Regex::new(
r"(?:^|\s)/b/(?P<id>\d+)(?:/(?P<mode>osu|taiko|fruits|mania))?(?:\+(?P<mods>[A-Z]+))?"
r"(?:^|\s|\W)(?P<main>/b/(?P<id>\d+)(?:/(?P<mode>osu|taiko|fruits|mania))?(?:\+(?P<mods>[A-Z]+))?)"
).unwrap();
}
pub fn hook(ctx: &mut Context, msg: &Message) -> () {
if msg.author.bot {
return;
}
let mut v = move || -> CommandResult {
let old_links = handle_old_links(ctx, &msg.content)?;
let new_links = handle_new_links(ctx, &msg.content)?;
let short_links = handle_short_links(ctx, &msg, &msg.content)?;
let mut last_beatmap = None;
for l in old_links
.into_iter()
.chain(new_links.into_iter())
.chain(short_links.into_iter())
{
if let Err(v) = msg.channel_id.send_message(&ctx, |m| match l.embed {
EmbedType::Beatmap(b, info, mods) => {
let t = handle_beatmap(&b, info, l.link, l.mode, mods, m);
let mode = l.mode.unwrap_or(b.mode);
last_beatmap = Some(super::BeatmapWithMode(b, mode));
t
}
EmbedType::Beatmapset(b) => handle_beatmapset(b, l.link, l.mode, m),
}) {
println!("Error in osu! hook: {:?}", v)
}
pub fn hook<'a>(
ctx: &'a Context,
msg: &'a Message,
) -> std::pin::Pin<Box<dyn future::Future<Output = Result<()>> + Send + 'a>> {
Box::pin(async move {
if msg.author.bot {
return Ok(());
}
let (old_links, new_links, short_links) = (
handle_old_links(ctx, &msg.content),
handle_new_links(ctx, &msg.content),
handle_short_links(ctx, &msg, &msg.content),
);
let last_beatmap = stream::select(old_links, stream::select(new_links, short_links))
.then(|l| async move {
let mut bm: Option<super::BeatmapWithMode> = None;
msg.channel_id
.send_message(&ctx, |m| match l.embed {
EmbedType::Beatmap(b, info, mods) => {
let t = handle_beatmap(&b, info, l.link, l.mode, mods, m);
let mode = l.mode.unwrap_or(b.mode);
bm = Some(super::BeatmapWithMode(b, mode));
t
}
EmbedType::Beatmapset(b) => handle_beatmapset(b, l.link, l.mode, m),
})
.await?;
let r: Result<_> = Ok(bm);
r
})
.filter_map(|v| async move {
match v {
Ok(v) => v,
Err(e) => {
eprintln!("{}", e);
None
}
}
})
.fold(None, |_, v| async move { Some(v) })
.await;
// Save the beatmap for query later.
if let Some(t) = last_beatmap {
if let Err(v) = super::cache::save_beatmap(&*ctx.data.read(), msg.channel_id, &t) {
dbg!(v);
}
super::cache::save_beatmap(&*ctx.data.read().await, msg.channel_id, &t)?;
}
Ok(())
};
if let Err(v) = v() {
println!("Error in osu! hook: {:?}", v)
}
})
}
enum EmbedType {
@ -80,167 +86,216 @@ struct ToPrint<'a> {
mode: Option<Mode>,
}
fn handle_old_links<'a>(ctx: &mut Context, content: &'a str) -> Result<Vec<ToPrint<'a>>, Error> {
let osu = ctx.data.get_cloned::<OsuClient>();
let mut to_prints: Vec<ToPrint<'a>> = Vec::new();
let cache = ctx.data.get_cloned::<BeatmapCache>();
for capture in OLD_LINK_REGEX.captures_iter(content) {
let req_type = capture.name("link_type").unwrap().as_str();
let req = match req_type {
"b" => BeatmapRequestKind::Beatmap(capture["id"].parse()?),
"s" => BeatmapRequestKind::Beatmapset(capture["id"].parse()?),
_ => continue,
};
let mode = capture
.name("mode")
.map(|v| v.as_str().parse())
.transpose()?
.and_then(|v| {
Some(match v {
0 => Mode::Std,
1 => Mode::Taiko,
2 => Mode::Catch,
3 => Mode::Mania,
_ => return None,
fn handle_old_links<'a>(
ctx: &'a Context,
content: &'a str,
) -> impl stream::Stream<Item = ToPrint<'a>> + 'a {
OLD_LINK_REGEX
.captures_iter(content)
.map(move |capture| async move {
let data = ctx.data.read().await;
let osu = data.get::<OsuClient>().unwrap();
let cache = data.get::<BeatmapCache>().unwrap();
let req_type = capture.name("link_type").unwrap().as_str();
let req = match req_type {
"b" => BeatmapRequestKind::Beatmap(capture["id"].parse()?),
"s" => BeatmapRequestKind::Beatmapset(capture["id"].parse()?),
_ => unreachable!(),
};
let mode = capture
.name("mode")
.map(|v| v.as_str().parse())
.transpose()?
.and_then(|v| {
Some(match v {
0 => Mode::Std,
1 => Mode::Taiko,
2 => Mode::Catch,
3 => Mode::Mania,
_ => return None,
})
});
let beatmaps = osu
.beatmaps(req, |v| match mode {
Some(m) => v.mode(m, true),
None => v,
})
});
let beatmaps = osu.beatmaps(req, |v| match mode {
Some(m) => v.mode(m, true),
None => v,
})?;
match req_type {
"b" => {
for b in beatmaps.into_iter() {
.await?;
if beatmaps.is_empty() {
return Ok(None);
}
let r: Result<_> = Ok(match req_type {
"b" => {
let b = beatmaps.into_iter().next().unwrap();
// collect beatmap info
let mods = capture
.name("mods")
.map(|v| Mods::from_str(v.as_str()).ok())
.flatten()
.unwrap_or(Mods::NOMOD);
let info = mode.unwrap_or(b.mode).to_oppai_mode().and_then(|mode| {
cache
let info = match mode.unwrap_or(b.mode).to_oppai_mode() {
Some(mode) => cache
.get_beatmap(b.beatmap_id)
.await
.and_then(|b| b.get_info_with(Some(mode), mods))
.ok()
});
to_prints.push(ToPrint {
.ok(),
None => None,
};
Some(ToPrint {
embed: EmbedType::Beatmap(b, info, mods),
link: capture.get(0).unwrap().as_str(),
mode,
})
}
}
"s" => to_prints.push(ToPrint {
embed: EmbedType::Beatmapset(beatmaps),
link: capture.get(0).unwrap().as_str(),
mode,
}),
_ => (),
}
}
Ok(to_prints)
"s" => Some(ToPrint {
embed: EmbedType::Beatmapset(beatmaps),
link: capture.get(0).unwrap().as_str(),
mode,
}),
_ => None,
});
r
})
.collect::<stream::FuturesUnordered<_>>()
.filter_map(|v| {
future::ready(match v {
Ok(v) => v,
Err(e) => {
eprintln!("{}", e);
None
}
})
})
}
fn handle_new_links<'a>(ctx: &mut Context, content: &'a str) -> Result<Vec<ToPrint<'a>>, Error> {
let osu = ctx.data.get_cloned::<OsuClient>();
let mut to_prints: Vec<ToPrint<'a>> = Vec::new();
let cache = ctx.data.get_cloned::<BeatmapCache>();
for capture in NEW_LINK_REGEX.captures_iter(content) {
let mode = capture
.name("mode")
.and_then(|v| Mode::parse_from_new_site(v.as_str()));
let link = capture.get(0).unwrap().as_str();
let req = match capture.name("beatmap_id") {
Some(ref v) => BeatmapRequestKind::Beatmap(v.as_str().parse()?),
None => {
BeatmapRequestKind::Beatmapset(capture.name("set_id").unwrap().as_str().parse()?)
fn handle_new_links<'a>(
ctx: &'a Context,
content: &'a str,
) -> impl stream::Stream<Item = ToPrint<'a>> + 'a {
NEW_LINK_REGEX
.captures_iter(content)
.map(|capture| async move {
let data = ctx.data.read().await;
let osu = data.get::<OsuClient>().unwrap();
let cache = data.get::<BeatmapCache>().unwrap();
let mode = capture
.name("mode")
.and_then(|v| Mode::parse_from_new_site(v.as_str()));
let link = capture.get(0).unwrap().as_str();
let req = match capture.name("beatmap_id") {
Some(ref v) => BeatmapRequestKind::Beatmap(v.as_str().parse()?),
None => BeatmapRequestKind::Beatmapset(
capture.name("set_id").unwrap().as_str().parse()?,
),
};
let beatmaps = osu
.beatmaps(req, |v| match mode {
Some(m) => v.mode(m, true),
None => v,
})
.await?;
if beatmaps.is_empty() {
return Ok(None);
}
};
let beatmaps = osu.beatmaps(req, |v| match mode {
Some(m) => v.mode(m, true),
None => v,
})?;
match capture.name("beatmap_id") {
Some(_) => {
for beatmap in beatmaps.into_iter() {
let r: Result<_> = Ok(match capture.name("beatmap_id") {
Some(_) => {
let beatmap = beatmaps.into_iter().next().unwrap();
// collect beatmap info
let mods = capture
.name("mods")
.and_then(|v| Mods::from_str(v.as_str()).ok())
.unwrap_or(Mods::NOMOD);
let info = mode
.unwrap_or(beatmap.mode)
.to_oppai_mode()
.and_then(|mode| {
cache
.get_beatmap(beatmap.beatmap_id)
.and_then(|b| b.get_info_with(Some(mode), mods))
.ok()
});
to_prints.push(ToPrint {
let info = match mode.unwrap_or(beatmap.mode).to_oppai_mode() {
Some(mode) => cache
.get_beatmap(beatmap.beatmap_id)
.await
.and_then(|b| b.get_info_with(Some(mode), mods))
.ok(),
None => None,
};
Some(ToPrint {
embed: EmbedType::Beatmap(beatmap, info, mods),
link,
mode,
})
}
}
None => to_prints.push(ToPrint {
embed: EmbedType::Beatmapset(beatmaps),
link,
mode,
}),
}
}
Ok(to_prints)
None => Some(ToPrint {
embed: EmbedType::Beatmapset(beatmaps),
link,
mode,
}),
});
r
})
.collect::<stream::FuturesUnordered<_>>()
.filter_map(|v| {
future::ready(match v {
Ok(v) => v,
Err(e) => {
eprintln!("{}", e);
None
}
})
})
}
fn handle_short_links<'a>(
ctx: &mut Context,
msg: &Message,
ctx: &'a Context,
msg: &'a Message,
content: &'a str,
) -> Result<Vec<ToPrint<'a>>, Error> {
if let Some(guild_id) = msg.guild_id {
if announcer::announcer_of(ctx, crate::discord::announcer::ANNOUNCER_KEY, guild_id)?
!= Some(msg.channel_id)
{
// Disable if we are not in the server's announcer channel
return Ok(vec![]);
}
}
let osu = ctx.data.get_cloned::<BeatmapMetaCache>();
let cache = ctx.data.get_cloned::<BeatmapCache>();
Ok(SHORT_LINK_REGEX
) -> impl stream::Stream<Item = ToPrint<'a>> + 'a {
SHORT_LINK_REGEX
.captures_iter(content)
.map(|capture| -> Result<_, Error> {
.map(|capture| async move {
if let Some(guild_id) = msg.guild_id {
if announcer::announcer_of(ctx, crate::discord::announcer::ANNOUNCER_KEY, guild_id)
.await?
!= Some(msg.channel_id)
{
// Disable if we are not in the server's announcer channel
return Err(Error::msg("not in server announcer channel"));
}
}
let data = ctx.data.read().await;
let osu = data.get::<BeatmapMetaCache>().unwrap();
let cache = data.get::<BeatmapCache>().unwrap();
let mode = capture
.name("mode")
.and_then(|v| Mode::parse_from_new_site(v.as_str()));
let id: u64 = capture.name("id").unwrap().as_str().parse()?;
let beatmap = match mode {
Some(mode) => osu.get_beatmap(id, mode),
None => osu.get_beatmap_default(id),
Some(mode) => osu.get_beatmap(id, mode).await,
None => osu.get_beatmap_default(id).await,
}?;
let mods = capture
.name("mods")
.and_then(|v| Mods::from_str(v.as_str()).ok())
.unwrap_or(Mods::NOMOD);
let info = mode
.unwrap_or(beatmap.mode)
.to_oppai_mode()
.and_then(|mode| {
cache
.get_beatmap(beatmap.beatmap_id)
.and_then(|b| b.get_info_with(Some(mode), mods))
.ok()
});
Ok(ToPrint {
let info = match mode.unwrap_or(beatmap.mode).to_oppai_mode() {
Some(mode) => cache
.get_beatmap(beatmap.beatmap_id)
.await
.and_then(|b| b.get_info_with(Some(mode), mods))
.ok(),
None => None,
};
let r: Result<_> = Ok(ToPrint {
embed: EmbedType::Beatmap(beatmap, info, mods),
link: capture.get(0).unwrap().as_str(),
link: capture.name("main").unwrap().as_str(),
mode,
});
r
})
.collect::<stream::FuturesUnordered<_>>()
.filter_map(|v| {
future::ready(match v {
Ok(v) => Some(v),
Err(e) => {
eprintln!("{}", e);
None
}
})
})
.filter_map(|v| v.ok())
.collect())
}
fn handle_beatmap<'a, 'b>(

556
youmubot-osu/src/discord/mod.rs
View file

@ -2,10 +2,9 @@ use crate::{
discord::beatmap_cache::BeatmapMetaCache,
discord::oppai_cache::BeatmapCache,
models::{Beatmap, Mode, Mods, Score, User},
request::{BeatmapRequestKind, UserID},
request::UserID,
Client as OsuHttpClient,
};
use rayon::prelude::*;
use serenity::{
framework::standard::{
macros::{command, group},
@ -14,7 +13,7 @@ use serenity::{
model::channel::Message,
utils::MessageBuilder,
};
use std::str::FromStr;
use std::{str::FromStr, sync::Arc};
use youmubot_prelude::*;
mod announcer;
@ -36,7 +35,7 @@ use server_rank::{LEADERBOARD_COMMAND, SERVER_RANK_COMMAND};
pub(crate) struct OsuClient;
impl TypeMapKey for OsuClient {
type Value = OsuHttpClient;
type Value = Arc<OsuHttpClient>;
}
/// Sets up the osu! command handling section.
@ -52,7 +51,7 @@ impl TypeMapKey for OsuClient {
///
pub fn setup(
path: &std::path::Path,
data: &mut ShareMap,
data: &mut TypeMap,
announcers: &mut AnnouncerHandler,
) -> CommandResult {
// Databases
@ -61,11 +60,10 @@ pub fn setup(
OsuUserBests::insert_into(&mut *data, &path.join("osu_user_bests.yaml"))?;
// API client
let http_client = data.get_cloned::<HTTPClient>();
let osu_client = OsuHttpClient::new(
http_client.clone(),
let http_client = data.get::<HTTPClient>().unwrap().clone();
let osu_client = Arc::new(OsuHttpClient::new(
std::env::var("OSU_API_KEY").expect("Please set OSU_API_KEY as osu! api key."),
);
));
data.insert::<OsuClient>(osu_client.clone());
data.insert::<oppai_cache::BeatmapCache>(oppai_cache::BeatmapCache::new(http_client));
data.insert::<beatmap_cache::BeatmapMetaCache>(beatmap_cache::BeatmapMetaCache::new(
@ -73,7 +71,7 @@ pub fn setup(
));
// Announcer
announcers.add(announcer::ANNOUNCER_KEY, announcer::updates);
announcers.add(announcer::ANNOUNCER_KEY, announcer::Announcer);
Ok(())
}
@ -101,8 +99,8 @@ struct Osu;
#[description = "Receive information about an user in osu!std mode."]
#[usage = "[username or user_id = your saved username]"]
#[max_args(1)]
pub fn std(ctx: &mut Context, msg: &Message, args: Args) -> CommandResult {
get_user(ctx, msg, args, Mode::Std)
pub async fn std(ctx: &Context, msg: &Message, args: Args) -> CommandResult {
get_user(ctx, msg, args, Mode::Std).await
}
#[command]
@ -110,8 +108,8 @@ pub fn std(ctx: &mut Context, msg: &Message, args: Args) -> CommandResult {
#[description = "Receive information about an user in osu!taiko mode."]
#[usage = "[username or user_id = your saved username]"]
#[max_args(1)]
pub fn taiko(ctx: &mut Context, msg: &Message, args: Args) -> CommandResult {
get_user(ctx, msg, args, Mode::Taiko)
pub async fn taiko(ctx: &Context, msg: &Message, args: Args) -> CommandResult {
get_user(ctx, msg, args, Mode::Taiko).await
}
#[command]
@ -119,8 +117,8 @@ pub fn taiko(ctx: &mut Context, msg: &Message, args: Args) -> CommandResult {
#[description = "Receive information about an user in osu!catch mode."]
#[usage = "[username or user_id = your saved username]"]
#[max_args(1)]
pub fn catch(ctx: &mut Context, msg: &Message, args: Args) -> CommandResult {
get_user(ctx, msg, args, Mode::Catch)
pub async fn catch(ctx: &Context, msg: &Message, args: Args) -> CommandResult {
get_user(ctx, msg, args, Mode::Catch).await
}
#[command]
@ -128,8 +126,8 @@ pub fn catch(ctx: &mut Context, msg: &Message, args: Args) -> CommandResult {
#[description = "Receive information about an user in osu!mania mode."]
#[usage = "[username or user_id = your saved username]"]
#[max_args(1)]
pub fn mania(ctx: &mut Context, msg: &Message, args: Args) -> CommandResult {
get_user(ctx, msg, args, Mode::Mania)
pub async fn mania(ctx: &Context, msg: &Message, args: Args) -> CommandResult {
get_user(ctx, msg, args, Mode::Mania).await
}
pub(crate) struct BeatmapWithMode(pub Beatmap, pub Mode);
@ -150,17 +148,15 @@ impl AsRef<Beatmap> for BeatmapWithMode {
#[description = "Save the given username as your username."]
#[usage = "[username or user_id]"]
#[num_args(1)]
pub fn save(ctx: &mut Context, msg: &Message, mut args: Args) -> CommandResult {
let osu = ctx.data.get_cloned::<OsuClient>();
pub async fn save(ctx: &Context, msg: &Message, mut args: Args) -> CommandResult {
let data = ctx.data.read().await;
let osu = data.get::<OsuClient>().unwrap();
let user = args.single::<String>()?;
let user: Option<User> = osu.user(UserID::Auto(user), |f| f)?;
let user: Option<User> = osu.user(UserID::Auto(user), |f| f).await?;
match user {
Some(u) => {
let db = OsuSavedUsers::open(&*ctx.data.read());
let mut db = db.borrow_mut()?;
db.insert(
OsuSavedUsers::open(&*data).borrow_mut()?.insert(
msg.author.id,
OsuUser {
id: u.id,
@ -174,10 +170,11 @@ pub fn save(ctx: &mut Context, msg: &Message, mut args: Args) -> CommandResult {
.push("user has been set to ")
.push_mono_safe(u.username)
.build(),
)?;
)
.await?;
}
None => {
msg.reply(&ctx, "user not found...")?;
msg.reply(&ctx, "user not found...").await?;
}
}
Ok(())
@ -200,7 +197,7 @@ impl FromStr for ModeArg {
fn to_user_id_query(
s: Option<UsernameArg>,
data: &ShareMap,
data: &TypeMap,
msg: &Message,
) -> Result<UserID, Error> {
let id = match s {
@ -236,151 +233,161 @@ impl FromStr for Nth {
}
}
fn list_plays(plays: Vec<Score>, mode: Mode, ctx: Context, m: &Message) -> CommandResult {
let watcher = ctx.data.get_cloned::<ReactionWatcher>();
let osu = ctx.data.get_cloned::<BeatmapMetaCache>();
let beatmap_cache = ctx.data.get_cloned::<BeatmapCache>();
async fn list_plays<'a>(
plays: Vec<Score>,
mode: Mode,
ctx: &'a Context,
m: &'a Message,
) -> CommandResult {
let plays = Arc::new(plays);
if plays.is_empty() {
m.reply(&ctx, "No plays found")?;
m.reply(&ctx, "No plays found").await?;
return Ok(());
}
let mut beatmaps: Vec<Option<String>> = vec![None; plays.len()];
const ITEMS_PER_PAGE: usize = 5;
let total_pages = (plays.len() + ITEMS_PER_PAGE - 1) / ITEMS_PER_PAGE;
watcher.paginate_fn(
ctx,
m.channel_id,
move |page, e| {
let page = page as usize;
let start = page * ITEMS_PER_PAGE;
let end = plays.len().min(start + ITEMS_PER_PAGE);
if start >= end {
return (e, Err(Error::from("No more pages")));
}
paginate(
move |page, ctx, msg| {
let plays = plays.clone();
Box::pin(async move {
let data = ctx.data.read().await;
let osu = data.get::<BeatmapMetaCache>().unwrap();
let beatmap_cache = data.get::<BeatmapCache>().unwrap();
let page = page as usize;
let start = page * ITEMS_PER_PAGE;
let end = plays.len().min(start + ITEMS_PER_PAGE);
if start >= end {
return Ok(false);
}
let plays = &plays[start..end];
let beatmaps: Vec<&mut String> = {
let b = &mut beatmaps[start..end];
b.par_iter_mut()
.enumerate()
.map(|(i, v)| {
v.get_or_insert_with(|| {
if let Some(b) = osu.get_beatmap(plays[i].beatmap_id, mode).ok() {
let stars = beatmap_cache
.get_beatmap(b.beatmap_id)
.ok()
.and_then(|b| {
mode.to_oppai_mode().and_then(|mode| {
b.get_info_with(Some(mode), plays[i].mods).ok()
})
})
.map(|info| info.stars as f64)
.unwrap_or(b.difficulty.stars);
format!(
"[{:.1}*] {} - {} [{}] ({})",
stars,
b.artist,
b.title,
b.difficulty_name,
b.short_link(Some(mode), Some(plays[i].mods)),
)
} else {
"FETCH_FAILED".to_owned()
}
})
let hourglass = msg.react(ctx, '⌛').await?;
let plays = &plays[start..end];
let beatmaps = plays
.iter()
.map(|play| async move {
let beatmap = osu.get_beatmap(play.beatmap_id, mode).await?;
let stars = {
let b = beatmap_cache.get_beatmap(beatmap.beatmap_id).await?;
mode.to_oppai_mode()
.and_then(|mode| b.get_info_with(Some(mode), play.mods).ok())
.map(|info| info.stars as f64)
.unwrap_or(beatmap.difficulty.stars)
};
let r: Result<_> = Ok(format!(
"[{:.1}*] {} - {} [{}] ({})",
stars,
beatmap.artist,
beatmap.title,
beatmap.difficulty_name,
beatmap.short_link(Some(mode), Some(play.mods)),
));
r
})
.collect::<Vec<_>>()
};
let pp = plays
.iter()
.map(|p| {
p.pp.map(|pp| format!("{:.2}pp", pp))
.or_else(|| {
beatmap_cache.get_beatmap(p.beatmap_id).ok().and_then(|b| {
mode.to_oppai_mode().and_then(|op| {
b.get_pp_from(
oppai_rs::Combo::NonFC {
max_combo: p.max_combo as u32,
misses: p.count_miss as u32,
},
p.accuracy(mode) as f32,
Some(op),
p.mods,
)
.ok()
.map(|pp| format!("{:.2}pp [?]", pp))
})
})
})
.unwrap_or("-".to_owned())
})
.collect::<Vec<_>>();
let pw = pp.iter().map(|v| v.len()).max().unwrap_or(2);
/*mods width*/
let mw = plays
.iter()
.map(|v| v.mods.to_string().len())
.max()
.unwrap()
.max(4);
/*beatmap names*/
let bw = beatmaps.iter().map(|v| v.len()).max().unwrap().max(7);
.collect::<stream::FuturesOrdered<_>>()
.map(|v| v.unwrap_or("FETCH_FAILED".to_owned()))
.collect::<Vec<String>>();
let pp = plays
.iter()
.map(|p| async move {
match p.pp.map(|pp| format!("{:.2}pp", pp)) {
Some(v) => Ok(v),
None => {
let b = beatmap_cache.get_beatmap(p.beatmap_id).await?;
let r: Result<_> = Ok(mode
.to_oppai_mode()
.and_then(|op| {
b.get_pp_from(
oppai_rs::Combo::NonFC {
max_combo: p.max_combo as u32,
misses: p.count_miss as u32,
},
p.accuracy(mode) as f32,
Some(op),
p.mods,
)
.ok()
.map(|pp| format!("{:.2}pp [?]", pp))
})
.unwrap_or("-".to_owned()));
r
}
}
})
.collect::<stream::FuturesOrdered<_>>()
.map(|v| v.unwrap_or("-".to_owned()))
.collect::<Vec<String>>();
let (beatmaps, pp) = future::join(beatmaps, pp).await;
let pw = pp.iter().map(|v| v.len()).max().unwrap_or(2);
/*mods width*/
let mw = plays
.iter()
.map(|v| v.mods.to_string().len())
.max()
.unwrap()
.max(4);
/*beatmap names*/
let bw = beatmaps.iter().map(|v| v.len()).max().unwrap().max(7);
let mut m = MessageBuilder::new();
// Table header
m.push_line(format!(
" # | {:pw$} | accuracy | rank | {:mw$} | {:bw$}",
"pp",
"mods",
"beatmap",
pw = pw,
mw = mw,
bw = bw
));
m.push_line(format!(
"------{:-<pw$}---------------------{:-<mw$}---{:-<bw$}",
"",
"",
"",
pw = pw,
mw = mw,
bw = bw
));
// Each row
for (id, (play, beatmap)) in plays.iter().zip(beatmaps.iter()).enumerate() {
let mut m = MessageBuilder::new();
// Table header
m.push_line(format!(
"{:>3} | {:>pw$} | {:>8} | {:^4} | {:mw$} | {:bw$}",
id + start + 1,
pp[id],
format!("{:.2}%", play.accuracy(mode)),
play.rank.to_string(),
play.mods.to_string(),
beatmap,
" # | {:pw$} | accuracy | rank | {:mw$} | {:bw$}",
"pp",
"mods",
"beatmap",
pw = pw,
mw = mw,
bw = bw
));
}
// End
let table = m.build().replace("```", "\\`\\`\\`");
let mut m = MessageBuilder::new();
m.push_codeblock(table, None).push_line(format!(
"Page **{}/{}**",
page + 1,
total_pages
));
if let None = mode.to_oppai_mode() {
m.push_line("Note: star difficulty doesn't reflect mods applied.");
} else {
m.push_line("[?] means pp was predicted by oppai-rs.");
}
(e.content(m.build()), Ok(()))
m.push_line(format!(
"------{:-<pw$}---------------------{:-<mw$}---{:-<bw$}",
"",
"",
"",
pw = pw,
mw = mw,
bw = bw
));
// Each row
for (id, (play, beatmap)) in plays.iter().zip(beatmaps.iter()).enumerate() {
m.push_line(format!(
"{:>3} | {:>pw$} | {:>8} | {:^4} | {:mw$} | {:bw$}",
id + start + 1,
pp[id],
format!("{:.2}%", play.accuracy(mode)),
play.rank.to_string(),
play.mods.to_string(),
beatmap,
pw = pw,
mw = mw,
bw = bw
));
}
// End
let table = m.build().replace("```", "\\`\\`\\`");
let mut m = MessageBuilder::new();
m.push_codeblock(table, None).push_line(format!(
"Page **{}/{}**",
page + 1,
total_pages
));
if let None = mode.to_oppai_mode() {
m.push_line("Note: star difficulty doesn't reflect mods applied.");
} else {
m.push_line("[?] means pp was predicted by oppai-rs.");
}
msg.edit(ctx, |f| f.content(m.to_string())).await?;
hourglass.delete(ctx).await?;
Ok(true)
})
},
ctx,
m.channel_id,
std::time::Duration::from_secs(60),
)
.await?;
Ok(())
}
#[command]
@ -388,44 +395,49 @@ fn list_plays(plays: Vec<Score>, mode: Mode, ctx: Context, m: &Message) -> Comma
#[usage = "#[the nth recent play = --all] / [mode (std, taiko, mania, catch) = std] / [username / user id = your saved id]"]
#[example = "#1 / taiko / natsukagami"]
#[max_args(3)]
pub fn recent(ctx: &mut Context, msg: &Message, mut args: Args) -> CommandResult {
pub async fn recent(ctx: &Context, msg: &Message, mut args: Args) -> CommandResult {
let data = ctx.data.read().await;
let nth = args.single::<Nth>().unwrap_or(Nth::All);
let mode = args.single::<ModeArg>().unwrap_or(ModeArg(Mode::Std)).0;
let user = to_user_id_query(args.single::<UsernameArg>().ok(), &*ctx.data.read(), msg)?;
let user = to_user_id_query(args.single::<UsernameArg>().ok(), &*data, msg)?;
let osu = ctx.data.get_cloned::<OsuClient>();
let meta_cache = ctx.data.get_cloned::<BeatmapMetaCache>();
let oppai = ctx.data.get_cloned::<BeatmapCache>();
let osu = data.get::<OsuClient>().unwrap();
let meta_cache = data.get::<BeatmapMetaCache>().unwrap();
let oppai = data.get::<BeatmapCache>().unwrap();
let user = osu
.user(user, |f| f.mode(mode))?
.user(user, |f| f.mode(mode))
.await?
.ok_or(Error::from("User not found"))?;
match nth {
Nth::Nth(nth) => {
let recent_play = osu
.user_recent(UserID::ID(user.id), |f| f.mode(mode).limit(nth))?
.user_recent(UserID::ID(user.id), |f| f.mode(mode).limit(nth))
.await?
.into_iter()
.last()
.ok_or(Error::from("No such play"))?;
let beatmap = meta_cache
.get_beatmap(recent_play.beatmap_id, mode)
.unwrap();
let content = oppai.get_beatmap(beatmap.beatmap_id)?;
let beatmap = meta_cache.get_beatmap(recent_play.beatmap_id, mode).await?;
let content = oppai.get_beatmap(beatmap.beatmap_id).await?;
let beatmap_mode = BeatmapWithMode(beatmap, mode);
msg.channel_id.send_message(&ctx, |m| {
m.content(format!(
"{}: here is the play that you requested",
msg.author
))
.embed(|m| score_embed(&recent_play, &beatmap_mode, &content, &user, None, m))
})?;
msg.channel_id
.send_message(&ctx, |m| {
m.content(format!(
"{}: here is the play that you requested",
msg.author
))
.embed(|m| score_embed(&recent_play, &beatmap_mode, &content, &user, None, m))
})
.await?;
// Save the beatmap...
cache::save_beatmap(&*ctx.data.read(), msg.channel_id, &beatmap_mode)?;
cache::save_beatmap(&*data, msg.channel_id, &beatmap_mode)?;
}
Nth::All => {
let plays = osu.user_recent(UserID::ID(user.id), |f| f.mode(mode).limit(50))?;
list_plays(plays, mode, ctx.clone(), msg)?;
let plays = osu
.user_recent(UserID::ID(user.id), |f| f.mode(mode).limit(50))
.await?;
list_plays(plays, mode, ctx, msg).await?;
}
}
Ok(())
@ -435,28 +447,33 @@ pub fn recent(ctx: &mut Context, msg: &Message, mut args: Args) -> CommandResult
#[description = "Show information from the last queried beatmap."]
#[usage = "[mods = no mod]"]
#[max_args(1)]
pub fn last(ctx: &mut Context, msg: &Message, mut args: Args) -> CommandResult {
let b = cache::get_beatmap(&*ctx.data.read(), msg.channel_id)?;
pub async fn last(ctx: &Context, msg: &Message, mut args: Args) -> CommandResult {
let data = ctx.data.read().await;
let b = cache::get_beatmap(&*data, msg.channel_id)?;
match b {
Some(BeatmapWithMode(b, m)) => {
let mods = args.find::<Mods>().unwrap_or(Mods::NOMOD);
let info = ctx
.data
.get_cloned::<BeatmapCache>()
.get_beatmap(b.beatmap_id)?
let info = data
.get::<BeatmapCache>()
.unwrap()
.get_beatmap(b.beatmap_id)
.await?
.get_info_with(m.to_oppai_mode(), mods)
.ok();
msg.channel_id.send_message(&ctx, |f| {
f.content(format!(
"{}: here is the beatmap you requested!",
msg.author
))
.embed(|c| beatmap_embed(&b, m, mods, info, c))
})?;
msg.channel_id
.send_message(&ctx, |f| {
f.content(format!(
"{}: here is the beatmap you requested!",
msg.author
))
.embed(|c| beatmap_embed(&b, m, mods, info, c))
})
.await?;
}
None => {
msg.reply(&ctx, "No beatmap was queried on this channel.")?;
msg.reply(&ctx, "No beatmap was queried on this channel.")
.await?;
}
}
@ -468,12 +485,14 @@ pub fn last(ctx: &mut Context, msg: &Message, mut args: Args) -> CommandResult {
#[usage = "[username or tag = yourself]"]
#[description = "Check your own or someone else's best record on the last beatmap. Also stores the result if possible."]
#[max_args(1)]
pub fn check(ctx: &mut Context, msg: &Message, mut args: Args) -> CommandResult {
let bm = cache::get_beatmap(&*ctx.data.read(), msg.channel_id)?;
pub async fn check(ctx: &Context, msg: &Message, mut args: Args) -> CommandResult {
let data = ctx.data.read().await;
let bm = cache::get_beatmap(&*data, msg.channel_id)?;
match bm {
None => {
msg.reply(&ctx, "No beatmap queried on this channel.")?;
msg.reply(&ctx, "No beatmap queried on this channel.")
.await?;
}
Some(bm) => {
let b = &bm.0;
@ -484,31 +503,36 @@ pub fn check(ctx: &mut Context, msg: &Message, mut args: Args) -> CommandResult
None => Some(msg.author.id),
_ => None,
};
let user = to_user_id_query(username_arg, &*ctx.data.read(), msg)?;
let user = to_user_id_query(username_arg, &*data, msg)?;
let osu = ctx.data.get_cloned::<OsuClient>();
let oppai = ctx.data.get_cloned::<BeatmapCache>();
let osu = data.get::<OsuClient>().unwrap();
let oppai = data.get::<BeatmapCache>().unwrap();
let content = oppai.get_beatmap(b.beatmap_id)?;
let content = oppai.get_beatmap(b.beatmap_id).await?;
let user = osu
.user(user, |f| f)?
.user(user, |f| f)
.await?
.ok_or(Error::from("User not found"))?;
let scores = osu.scores(b.beatmap_id, |f| f.user(UserID::ID(user.id)).mode(m))?;
let scores = osu
.scores(b.beatmap_id, |f| f.user(UserID::ID(user.id)).mode(m))
.await?;
if scores.is_empty() {
msg.reply(&ctx, "No scores found")?;
msg.reply(&ctx, "No scores found").await?;
}
for score in scores.iter() {
msg.channel_id.send_message(&ctx, |c| {
c.embed(|m| score_embed(score, &bm, &content, &user, None, m))
})?;
msg.channel_id
.send_message(&ctx, |c| {
c.embed(|m| score_embed(&score, &bm, &content, &user, None, m))
})
.await?;
}
if let Some(user_id) = user_id {
// Save to database
OsuUserBests::open(&*ctx.data.read())
OsuUserBests::open(&*data)
.borrow_mut()?
.entry((bm.0.beatmap_id, bm.1))
.or_default()
@ -522,27 +546,32 @@ pub fn check(ctx: &mut Context, msg: &Message, mut args: Args) -> CommandResult
#[command]
#[description = "Get the n-th top record of an user."]
#[usage = "#[n-th = --all] / [mode (std, taiko, catch, mania)] = std / [username or user_id = your saved user id]"]
#[example = "#2 / taiko / natsukagami"]
#[usage = "[mode (std, taiko, catch, mania)] = std / #[n-th = --all] / [username or user_id = your saved user id]"]
#[example = "taiko / #2 / natsukagami"]
#[max_args(3)]
pub fn top(ctx: &mut Context, msg: &Message, mut args: Args) -> CommandResult {
pub async fn top(ctx: &Context, msg: &Message, mut args: Args) -> CommandResult {
let data = ctx.data.read().await;
let nth = args.single::<Nth>().unwrap_or(Nth::All);
let mode = args
.single::<ModeArg>()
.map(|ModeArg(t)| t)
.unwrap_or(Mode::Std);
let user = to_user_id_query(args.single::<UsernameArg>().ok(), &*ctx.data.read(), msg)?;
let user = to_user_id_query(args.single::<UsernameArg>().ok(), &*data, msg)?;
let meta_cache = data.get::<BeatmapMetaCache>().unwrap();
let osu = data.get::<OsuClient>().unwrap();
let osu = ctx.data.get_cloned::<OsuClient>();
let oppai = ctx.data.get_cloned::<BeatmapCache>();
let oppai = data.get::<BeatmapCache>().unwrap();
let user = osu
.user(user, |f| f.mode(mode))?
.user(user, |f| f.mode(mode))
.await?
.ok_or(Error::from("User not found"))?;
match nth {
Nth::Nth(nth) => {
let top_play = osu.user_best(UserID::ID(user.id), |f| f.mode(mode).limit(nth))?;
let top_play = osu
.user_best(UserID::ID(user.id), |f| f.mode(mode).limit(nth))
.await?;
let rank = top_play.len() as u8;
@ -550,69 +579,76 @@ pub fn top(ctx: &mut Context, msg: &Message, mut args: Args) -> CommandResult {
.into_iter()
.last()
.ok_or(Error::from("No such play"))?;
let beatmap = osu
.beatmaps(BeatmapRequestKind::Beatmap(top_play.beatmap_id), |f| {
f.mode(mode, true)
})?
.into_iter()
.next()
.unwrap();
let content = oppai.get_beatmap(beatmap.beatmap_id)?;
let beatmap = meta_cache.get_beatmap(top_play.beatmap_id, mode).await?;
let content = oppai.get_beatmap(beatmap.beatmap_id).await?;
let beatmap = BeatmapWithMode(beatmap, mode);
msg.channel_id.send_message(&ctx, |m| {
m.content(format!(
"{}: here is the play that you requested",
msg.author
))
.embed(|m| score_embed(&top_play, &beatmap, &content, &user, Some(rank), m))
})?;
msg.channel_id
.send_message(&ctx, |m| {
m.content(format!(
"{}: here is the play that you requested",
msg.author
))
.embed(|m| score_embed(&top_play, &beatmap, &content, &user, Some(rank), m))
})
.await?;
// Save the beatmap...
cache::save_beatmap(&*ctx.data.read(), msg.channel_id, &beatmap)?;
cache::save_beatmap(&*data, msg.channel_id, &beatmap)?;
}
Nth::All => {
let plays = osu.user_best(UserID::ID(user.id), |f| f.mode(mode).limit(100))?;
list_plays(plays, mode, ctx.clone(), msg)?;
let plays = osu
.user_best(UserID::ID(user.id), |f| f.mode(mode).limit(100))
.await?;
list_plays(plays, mode, ctx, msg).await?;
}
}
Ok(())
}
fn get_user(ctx: &mut Context, msg: &Message, mut args: Args, mode: Mode) -> CommandResult {
let user = to_user_id_query(args.single::<UsernameArg>().ok(), &*ctx.data.read(), msg)?;
let osu = ctx.data.get_cloned::<OsuClient>();
let cache = ctx.data.get_cloned::<BeatmapMetaCache>();
let user = osu.user(user, |f| f.mode(mode))?;
let oppai = ctx.data.get_cloned::<BeatmapCache>();
async fn get_user(ctx: &Context, msg: &Message, mut args: Args, mode: Mode) -> CommandResult {
let data = ctx.data.read().await;
let user = to_user_id_query(args.single::<UsernameArg>().ok(), &*data, msg)?;
let osu = data.get::<OsuClient>().unwrap();
let cache = data.get::<BeatmapMetaCache>().unwrap();
let user = osu.user(user, |f| f.mode(mode)).await?;
let oppai = data.get::<BeatmapCache>().unwrap();
match user {
Some(u) => {
let best = osu
.user_best(UserID::ID(u.id), |f| f.limit(1).mode(mode))?
let best = match osu
.user_best(UserID::ID(u.id), |f| f.limit(1).mode(mode))
.await?
.into_iter()
.next()
.map(|m| -> Result<_, Error> {
let beatmap = cache.get_beatmap(m.beatmap_id, mode)?;
let info = mode
.to_oppai_mode()
.map(|mode| -> Result<_, Error> {
Ok(oppai
.get_beatmap(m.beatmap_id)?
.get_info_with(Some(mode), m.mods)?)
})
.transpose()?;
Ok((m, BeatmapWithMode(beatmap, mode), info))
{
Some(m) => {
let beatmap = cache.get_beatmap(m.beatmap_id, mode).await?;
let info = match mode.to_oppai_mode() {
Some(mode) => Some(
oppai
.get_beatmap(m.beatmap_id)
.await?
.get_info_with(Some(mode), m.mods)?,
),
None => None,
};
Some((m, BeatmapWithMode(beatmap, mode), info))
}
None => None,
};
msg.channel_id
.send_message(&ctx, |m| {
m.content(format!(
"{}: here is the user that you requested",
msg.author
))
.embed(|m| user_embed(u, best, m))
})
.transpose()?;
msg.channel_id.send_message(&ctx, |m| {
m.content(format!(
"{}: here is the user that you requested",
msg.author
))
.embed(|m| user_embed(u, best, m))
})
.await?;
}
None => msg.reply(&ctx, "🔍 user not found!"),
}?;
None => {
msg.reply(&ctx, "🔍 user not found!").await?;
}
};
Ok(())
}

47
youmubot-osu/src/discord/oppai_cache.rs
View file

@ -1,12 +1,11 @@
use serenity::framework::standard::CommandError;
use std::{ffi::CString, sync::Arc};
use youmubot_prelude::TypeMapKey;
use youmubot_prelude::*;
/// the information collected from a download/Oppai request.
#[derive(Clone, Debug)]
#[derive(Debug)]
pub struct BeatmapContent {
id: u64,
content: Arc<CString>,
content: CString,
}
/// the output of "one" oppai run.
@ -24,7 +23,7 @@ impl BeatmapContent {
accuracy: f32,
mode: Option<oppai_rs::Mode>,
mods: impl Into<oppai_rs::Mods>,
) -> Result<f32, CommandError> {
) -> Result<f32> {
let mut oppai = oppai_rs::Oppai::new_from_content(&self.content[..])?;
oppai.combo(combo)?.accuracy(accuracy)?.mods(mods.into());
if let Some(mode) = mode {
@ -38,7 +37,7 @@ impl BeatmapContent {
&self,
mode: Option<oppai_rs::Mode>,
mods: impl Into<oppai_rs::Mods>,
) -> Result<BeatmapInfo, CommandError> {
) -> Result<BeatmapInfo> {
let mut oppai = oppai_rs::Oppai::new_from_content(&self.content[..])?;
if let Some(mode) = mode {
oppai.mode(mode)?;
@ -56,39 +55,47 @@ impl BeatmapContent {
}
/// A central cache for the beatmaps.
#[derive(Clone, Debug)]
pub struct BeatmapCache {
client: reqwest::blocking::Client,
cache: Arc<dashmap::DashMap<u64, BeatmapContent>>,
client: ratelimit::Ratelimit<reqwest::Client>,
cache: dashmap::DashMap<u64, Arc<BeatmapContent>>,
}
impl BeatmapCache {
/// Create a new cache.
pub fn new(client: reqwest::blocking::Client) -> Self {
pub fn new(client: reqwest::Client) -> Self {
let client = ratelimit::Ratelimit::new(client, 5, std::time::Duration::from_secs(1));
BeatmapCache {
client,
cache: Arc::new(dashmap::DashMap::new()),
cache: dashmap::DashMap::new(),
}
}
fn download_beatmap(&self, id: u64) -> Result<BeatmapContent, CommandError> {
async fn download_beatmap(&self, id: u64) -> Result<BeatmapContent> {
let content = self
.client
.borrow()
.await?
.get(&format!("https://osu.ppy.sh/osu/{}", id))
.send()?
.bytes()?;
.send()
.await?
.bytes()
.await?;
Ok(BeatmapContent {
id,
content: Arc::new(CString::new(content.into_iter().collect::<Vec<_>>())?),
content: CString::new(content.into_iter().collect::<Vec<_>>())?,
})
}
/// Get a beatmap from the cache.
pub fn get_beatmap(&self, id: u64) -> Result<BeatmapContent, CommandError> {
self.cache
.entry(id)
.or_try_insert_with(|| self.download_beatmap(id))
.map(|v| v.clone())
pub async fn get_beatmap(
&self,
id: u64,
) -> Result<impl std::ops::Deref<Target = BeatmapContent>> {
if !self.cache.contains_key(&id) {
self.cache
.insert(id, Arc::new(self.download_beatmap(id).await?));
}
Ok(self.cache.get(&id).unwrap().clone())
}
}

378
youmubot-osu/src/discord/server_rank.rs
View file

@ -1,12 +1,14 @@
use super::{
cache::get_beatmap,
db::{OsuSavedUsers, OsuUserBests},
ModeArg,
ModeArg, OsuClient,
};
use crate::{
models::{Mode, Score},
request::UserID,
};
use crate::models::{Mode, Score};
use serenity::{
builder::EditMessage,
framework::standard::{macros::command, Args, CommandError as Error, CommandResult},
framework::standard::{macros::command, Args, CommandResult},
model::channel::Message,
utils::MessageBuilder,
};
@ -17,15 +19,15 @@ use youmubot_prelude::*;
#[usage = "[mode (Std, Taiko, Catch, Mania) = Std]"]
#[max_args(1)]
#[only_in(guilds)]
pub fn server_rank(ctx: &mut Context, m: &Message, mut args: Args) -> CommandResult {
pub async fn server_rank(ctx: &Context, m: &Message, mut args: Args) -> CommandResult {
let data = ctx.data.read().await;
let mode = args.single::<ModeArg>().map(|v| v.0).unwrap_or(Mode::Std);
let guild = m.guild_id.expect("Guild-only command");
let users = OsuSavedUsers::open(&*ctx.data.read())
.borrow()
.expect("DB initialized")
.iter()
.filter_map(|(user_id, osu_user)| {
guild.member(&ctx, user_id).ok().and_then(|member| {
let users = OsuSavedUsers::open(&*data).borrow()?.clone();
let users = users
.into_iter()
.map(|(user_id, osu_user)| async move {
guild.member(&ctx, user_id).await.ok().and_then(|member| {
osu_user
.pp
.get(mode as usize)
@ -34,7 +36,10 @@ pub fn server_rank(ctx: &mut Context, m: &Message, mut args: Args) -> CommandRes
.map(|pp| (pp, member.distinct(), osu_user.last_update.clone()))
})
})
.collect::<Vec<_>>();
.collect::<stream::FuturesUnordered<_>>()
.filter_map(|v| future::ready(v))
.collect::<Vec<_>>()
.await;
let last_update = users.iter().map(|(_, _, a)| a).min().cloned();
let mut users = users
.into_iter()
@ -43,47 +48,55 @@ pub fn server_rank(ctx: &mut Context, m: &Message, mut args: Args) -> CommandRes
users.sort_by(|(a, _), (b, _)| (*b).partial_cmp(a).unwrap_or(std::cmp::Ordering::Equal));
if users.is_empty() {
m.reply(&ctx, "No saved users in the current server...")?;
m.reply(&ctx, "No saved users in the current server...")
.await?;
return Ok(());
}
let users = std::sync::Arc::new(users);
let last_update = last_update.unwrap();
const ITEMS_PER_PAGE: usize = 10;
ctx.data.get_cloned::<ReactionWatcher>().paginate_fn(
ctx.clone(),
m.channel_id,
move |page: u8, e: &mut EditMessage| {
let start = (page as usize) * ITEMS_PER_PAGE;
let end = (start + ITEMS_PER_PAGE).min(users.len());
if start >= end {
return (e, Err(Error("No more items".to_owned())));
}
let total_len = users.len();
let users = &users[start..end];
let username_len = users.iter().map(|(_, u)| u.len()).max().unwrap().max(8);
let mut content = MessageBuilder::new();
content
.push_line("```")
.push_line("Rank | pp | Username")
.push_line(format!("-----------------{:-<uw$}", "", uw = username_len));
for (id, (pp, member)) in users.iter().enumerate() {
paginate(
move |page: u8, ctx: &Context, m: &mut Message| {
const ITEMS_PER_PAGE: usize = 10;
let users = users.clone();
Box::pin(async move {
let start = (page as usize) * ITEMS_PER_PAGE;
let end = (start + ITEMS_PER_PAGE).min(users.len());
if start >= end {
return Ok(false);
}
let total_len = users.len();
let users = &users[start..end];
let username_len = users.iter().map(|(_, u)| u.len()).max().unwrap_or(8).max(8);
let mut content = MessageBuilder::new();
content
.push(format!(
"{:>4} | {:>7.2} | ",
format!("#{}", 1 + id + start),
pp
))
.push_line_safe(member);
}
content.push_line("```").push_line(format!(
"Page **{}**/**{}**. Last updated: `{}`",
page + 1,
(total_len + ITEMS_PER_PAGE - 1) / ITEMS_PER_PAGE,
last_update.to_rfc2822()
));
(e.content(content.build()), Ok(()))
.push_line("```")
.push_line("Rank | pp | Username")
.push_line(format!("-----------------{:-<uw$}", "", uw = username_len));
for (id, (pp, member)) in users.iter().enumerate() {
content
.push(format!(
"{:>4} | {:>7.2} | ",
format!("#{}", 1 + id + start),
pp
))
.push_line_safe(member);
}
content.push_line("```").push_line(format!(
"Page **{}**/**{}**. Last updated: `{}`",
page + 1,
(total_len + ITEMS_PER_PAGE - 1) / ITEMS_PER_PAGE,
last_update.to_rfc2822()
));
m.edit(ctx, |f| f.content(content.to_string())).await?;
Ok(true)
})
},
ctx,
m.channel_id,
std::time::Duration::from_secs(60),
)?;
)
.await?;
Ok(())
}
@ -93,48 +106,79 @@ pub fn server_rank(ctx: &mut Context, m: &Message, mut args: Args) -> CommandRes
#[description = "See the server's ranks on the last seen beatmap"]
#[max_args(0)]
#[only_in(guilds)]
pub fn leaderboard(ctx: &mut Context, m: &Message, mut _args: Args) -> CommandResult {
let bm = match get_beatmap(&*ctx.data.read(), m.channel_id)? {
pub async fn leaderboard(ctx: &Context, m: &Message, mut _args: Args) -> CommandResult {
let data = ctx.data.read().await;
let mut osu_user_bests = OsuUserBests::open(&*data);
let bm = match get_beatmap(&*data, m.channel_id)? {
Some(bm) => bm,
None => {
m.reply(&ctx, "No beatmap queried on this channel.")?;
m.reply(&ctx, "No beatmap queried on this channel.").await?;
return Ok(());
}
};
// Run a check on the user once too!
{
let osu_users = OsuSavedUsers::open(&*data);
let user = osu_users.borrow()?.get(&m.author.id).map(|v| v.id);
if let Some(id) = user {
let osu = data.get::<OsuClient>().unwrap();
if let Ok(scores) = osu
.scores(bm.0.beatmap_id, |f| f.user(UserID::ID(id)))
.await
{
if !scores.is_empty() {
osu_user_bests
.borrow_mut()?
.entry((bm.0.beatmap_id, bm.1))
.or_default()
.insert(m.author.id, scores);
}
}
}
}
let guild = m.guild_id.expect("Guild-only command");
let scores = {
let users = OsuUserBests::open(&*ctx.data.read());
let users = users.borrow()?;
let users = match users.get(&(bm.0.beatmap_id, bm.1)) {
const NO_SCORES: &'static str =
"No scores have been recorded for this beatmap. Run `osu check` to scan for yours!";
let users = osu_user_bests
.borrow()?
.get(&(bm.0.beatmap_id, bm.1))
.cloned();
let users = match users {
None => {
m.reply(
&ctx,
"No scores have been recorded for this beatmap. Run `osu check` to scan for yours!",
)?;
m.reply(&ctx, NO_SCORES).await?;
return Ok(());
}
Some(v) if v.is_empty() => {
m.reply(
&ctx,
"No scores have been recorded for this beatmap. Run `osu check` to scan for yours!",
)?;
m.reply(&ctx, NO_SCORES).await?;
return Ok(());
}
Some(v) => v,
};
let mut scores: Vec<(f64, String, Score)> = users
.iter()
.filter_map(|(user_id, scores)| {
.into_iter()
.map(|(user_id, scores)| async move {
guild
.member(&ctx, user_id)
.await
.ok()
.and_then(|m| Some((m.distinct(), scores)))
})
.flat_map(|(user, scores)| scores.into_iter().map(move |v| (user.clone(), v.clone())))
.filter_map(|(user, score)| score.pp.map(|v| (v, user, score)))
.collect::<Vec<_>>();
.collect::<stream::FuturesUnordered<_>>()
.filter_map(|v| future::ready(v))
.flat_map(|(user, scores)| {
scores
.into_iter()
.map(move |v| future::ready((user.clone(), v.clone())))
.collect::<stream::FuturesUnordered<_>>()
})
.filter_map(|(user, score)| future::ready(score.pp.map(|v| (v, user, score))))
.collect::<Vec<_>>()
.await;
scores
.sort_by(|(a, _, _), (b, _, _)| b.partial_cmp(a).unwrap_or(std::cmp::Ordering::Equal));
scores
@ -144,115 +188,121 @@ pub fn leaderboard(ctx: &mut Context, m: &Message, mut _args: Args) -> CommandRe
m.reply(
&ctx,
"No scores have been recorded for this beatmap. Run `osu check` to scan for yours!",
)?;
)
.await?;
return Ok(());
}
ctx.data.get_cloned::<ReactionWatcher>().paginate_fn(
ctx.clone(),
m.channel_id,
move |page: u8, e: &mut EditMessage| {
paginate(
move |page: u8, ctx: &Context, m: &mut Message| {
const ITEMS_PER_PAGE: usize = 5;
let start = (page as usize) * ITEMS_PER_PAGE;
let end = (start + ITEMS_PER_PAGE).min(scores.len());
if start >= end {
return (e, Err(Error("No more items".to_owned())));
return Box::pin(future::ready(Ok(false)));
}
let total_len = scores.len();
let scores = &scores[start..end];
// username width
let uw = scores
.iter()
.map(|(_, u, _)| u.len())
.max()
.unwrap_or(8)
.max(8);
let accuracies = scores
.iter()
.map(|(_, _, v)| format!("{:.2}%", v.accuracy(bm.1)))
.collect::<Vec<_>>();
let aw = accuracies.iter().map(|v| v.len()).max().unwrap().max(3);
let misses = scores
.iter()
.map(|(_, _, v)| format!("{}", v.count_miss))
.collect::<Vec<_>>();
let mw = misses.iter().map(|v| v.len()).max().unwrap().max(4);
let ranks = scores
.iter()
.map(|(_, _, v)| v.rank.to_string())
.collect::<Vec<_>>();
let rw = ranks.iter().map(|v| v.len()).max().unwrap().max(4);
let pp = scores
.iter()
.map(|(pp, _, _)| format!("{:.2}", pp))
.collect::<Vec<_>>();
let pw = pp.iter().map(|v| v.len()).max().unwrap_or(2);
/*mods width*/
let mdw = scores
.iter()
.map(|(_, _, v)| v.mods.to_string().len())
.max()
.unwrap()
.max(4);
let mut content = MessageBuilder::new();
content
.push_line("```")
.push_line(format!(
"rank | {:>pw$} | {:mdw$} | {:rw$} | {:>aw$} | {:mw$} | {:uw$}",
"pp",
"mods",
"rank",
"acc",
"miss",
"user",
pw = pw,
mdw = mdw,
rw = rw,
aw = aw,
mw = mw,
uw = uw,
))
.push_line(format!(
"-------{:-<pw$}---{:-<mdw$}---{:-<rw$}---{:-<aw$}---{:-<mw$}---{:-<uw$}",
"",
"",
"",
"",
"",
"",
pw = pw,
mdw = mdw,
rw = rw,
aw = aw,
mw = mw,
uw = uw,
let scores = (&scores[start..end]).iter().cloned().collect::<Vec<_>>();
let bm = (bm.0.clone(), bm.1.clone());
Box::pin(async move {
// username width
let uw = scores
.iter()
.map(|(_, u, _)| u.len())
.max()
.unwrap_or(8)
.max(8);
let accuracies = scores
.iter()
.map(|(_, _, v)| format!("{:.2}%", v.accuracy(bm.1)))
.collect::<Vec<_>>();
let aw = accuracies.iter().map(|v| v.len()).max().unwrap().max(3);
let misses = scores
.iter()
.map(|(_, _, v)| format!("{}", v.count_miss))
.collect::<Vec<_>>();
let mw = misses.iter().map(|v| v.len()).max().unwrap().max(4);
let ranks = scores
.iter()
.map(|(_, _, v)| v.rank.to_string())
.collect::<Vec<_>>();
let rw = ranks.iter().map(|v| v.len()).max().unwrap().max(4);
let pp = scores
.iter()
.map(|(pp, _, _)| format!("{:.2}", pp))
.collect::<Vec<_>>();
let pw = pp.iter().map(|v| v.len()).max().unwrap_or(2);
/*mods width*/
let mdw = scores
.iter()
.map(|(_, _, v)| v.mods.to_string().len())
.max()
.unwrap()
.max(4);
let mut content = MessageBuilder::new();
content
.push_line("```")
.push_line(format!(
"rank | {:>pw$} | {:mdw$} | {:rw$} | {:>aw$} | {:mw$} | {:uw$}",
"pp",
"mods",
"rank",
"acc",
"miss",
"user",
pw = pw,
mdw = mdw,
rw = rw,
aw = aw,
mw = mw,
uw = uw,
))
.push_line(format!(
"-------{:-<pw$}---{:-<mdw$}---{:-<rw$}---{:-<aw$}---{:-<mw$}---{:-<uw$}",
"",
"",
"",
"",
"",
"",
pw = pw,
mdw = mdw,
rw = rw,
aw = aw,
mw = mw,
uw = uw,
));
for (id, (_, member, p)) in scores.iter().enumerate() {
content.push_line_safe(format!(
"{:>4} | {:>pw$} | {:>mdw$} | {:>rw$} | {:>aw$} | {:>mw$} | {:uw$}",
format!("#{}", 1 + id + start),
pp[id],
p.mods.to_string(),
ranks[id],
accuracies[id],
misses[id],
member,
pw = pw,
mdw = mdw,
rw = rw,
aw = aw,
mw = mw,
uw = uw,
));
}
content.push_line("```").push_line(format!(
"Page **{}**/**{}**. Not seeing your scores? Run `osu check` to update.",
page + 1,
(total_len + ITEMS_PER_PAGE - 1) / ITEMS_PER_PAGE,
));
for (id, (_, member, p)) in scores.iter().enumerate() {
content.push_line_safe(format!(
"{:>4} | {:>pw$} | {:>mdw$} | {:>rw$} | {:>aw$} | {:>mw$} | {:uw$}",
format!("#{}", 1 + id + start),
pp[id],
p.mods.to_string(),
ranks[id],
accuracies[id],
misses[id],
member,
pw = pw,
mdw = mdw,
rw = rw,
aw = aw,
mw = mw,
uw = uw,
));
}
content.push_line("```").push_line(format!(
"Page **{}**/**{}**. Not seeing your scores? Run `osu check` to update.",
page + 1,
(total_len + ITEMS_PER_PAGE - 1) / ITEMS_PER_PAGE,
));
(e.content(content.build()), Ok(()))
m.edit(&ctx, |f| f.content(content.build())).await?;
Ok(true)
})
},
ctx,
m.channel_id,
std::time::Duration::from_secs(60),
)?;
)
.await?;
Ok(())
}

64
youmubot-osu/src/lib.rs
View file

@ -8,16 +8,17 @@ mod test;
use models::*;
use request::builders::*;
use request::*;
use reqwest::blocking::{Client as HTTPClient, RequestBuilder, Response};
use serenity::framework::standard::CommandError as Error;
use std::{convert::TryInto, sync::Arc};
use reqwest::Client as HTTPClient;
use std::convert::TryInto;
use youmubot_prelude::{ratelimit::Ratelimit, *};
/// The number of requests per minute to the osu! server.
const REQUESTS_PER_MINUTE: usize = 200;
/// Client is the client that will perform calls to the osu! api server.
/// It's cheap to clone, so do it.
#[derive(Clone, Debug)]
pub struct Client {
key: Arc<String>,
client: HTTPClient,
client: Ratelimit<HTTPClient>,
key: String,
}
fn vec_try_into<U, T: std::convert::TryFrom<U>>(v: Vec<U>) -> Result<Vec<T>, T::Error> {
@ -32,50 +33,55 @@ fn vec_try_into<U, T: std::convert::TryFrom<U>>(v: Vec<U>) -> Result<Vec<T>, T::
impl Client {
/// Create a new client from the given API key.
pub fn new(http_client: HTTPClient, key: String) -> Client {
Client {
key: Arc::new(key),
client: http_client,
}
pub fn new(key: String) -> Client {
let client = Ratelimit::new(
HTTPClient::new(),
REQUESTS_PER_MINUTE,
std::time::Duration::from_secs(60),
);
Client { key, client }
}
fn build_request(&self, r: RequestBuilder) -> Result<Response, Error> {
let v = r.query(&[("k", &*self.key)]).build()?;
// dbg!(v.url());
Ok(self.client.execute(v)?)
pub(crate) async fn build_request(&self, url: &str) -> Result<reqwest::RequestBuilder> {
Ok(self
.client
.borrow()
.await?
.get(url)
.query(&[("k", &*self.key)]))
}
pub fn beatmaps(
pub async fn beatmaps(
&self,
kind: BeatmapRequestKind,
f: impl FnOnce(&mut BeatmapRequestBuilder) -> &mut BeatmapRequestBuilder,
) -> Result<Vec<Beatmap>, Error> {
) -> Result<Vec<Beatmap>> {
let mut r = BeatmapRequestBuilder::new(kind);
f(&mut r);
let res: Vec<raw::Beatmap> = self.build_request(r.build(&self.client))?.json()?;
let res: Vec<raw::Beatmap> = r.build(&self).await?.json().await?;
Ok(vec_try_into(res)?)
}
pub fn user(
pub async fn user(
&self,
user: UserID,
f: impl FnOnce(&mut UserRequestBuilder) -> &mut UserRequestBuilder,
) -> Result<Option<User>, Error> {
let mut r = UserRequestBuilder::new(user);
f(&mut r);
let res: Vec<raw::User> = self.build_request(r.build(&self.client))?.json()?;
let res: Vec<raw::User> = r.build(&self).await?.json().await?;
let res = vec_try_into(res)?;
Ok(res.into_iter().next())
}
pub fn scores(
pub async fn scores(
&self,
beatmap_id: u64,
f: impl FnOnce(&mut ScoreRequestBuilder) -> &mut ScoreRequestBuilder,
) -> Result<Vec<Score>, Error> {
let mut r = ScoreRequestBuilder::new(beatmap_id);
f(&mut r);
let res: Vec<raw::Score> = self.build_request(r.build(&self.client))?.json()?;
let res: Vec<raw::Score> = r.build(&self).await?.json().await?;
let mut res: Vec<Score> = vec_try_into(res)?;
// with a scores request you need to fill the beatmap ids yourself
@ -85,23 +91,23 @@ impl Client {
Ok(res)
}
pub fn user_best(
pub async fn user_best(
&self,
user: UserID,
f: impl FnOnce(&mut UserScoreRequestBuilder) -> &mut UserScoreRequestBuilder,
) -> Result<Vec<Score>, Error> {
self.user_scores(UserScoreType::Best, user, f)
self.user_scores(UserScoreType::Best, user, f).await
}
pub fn user_recent(
pub async fn user_recent(
&self,
user: UserID,
f: impl FnOnce(&mut UserScoreRequestBuilder) -> &mut UserScoreRequestBuilder,
) -> Result<Vec<Score>, Error> {
self.user_scores(UserScoreType::Recent, user, f)
self.user_scores(UserScoreType::Recent, user, f).await
}
fn user_scores(
async fn user_scores(
&self,
u: UserScoreType,
user: UserID,
@ -109,7 +115,7 @@ impl Client {
) -> Result<Vec<Score>, Error> {
let mut r = UserScoreRequestBuilder::new(u, user);
f(&mut r);
let res: Vec<raw::Score> = self.build_request(r.build(&self.client))?.json()?;
let res: Vec<raw::Score> = r.build(&self).await?.json().await?;
let res = vec_try_into(res)?;
Ok(res)
}

41
youmubot-osu/src/request.rs
View file

@ -1,6 +1,7 @@
use crate::models::{Mode, Mods};
use crate::Client;
use chrono::{DateTime, Utc};
use reqwest::blocking::{Client, RequestBuilder};
use youmubot_prelude::*;
trait ToQuery {
fn to_query(&self) -> Vec<(&'static str, String)>;
@ -84,6 +85,8 @@ impl ToQuery for BeatmapRequestKind {
}
pub mod builders {
use reqwest::Response;
use super::*;
/// A builder for a Beatmap request.
pub struct BeatmapRequestBuilder {
@ -110,12 +113,15 @@ pub mod builders {
self
}
pub(crate) fn build(self, client: &Client) -> RequestBuilder {
client
.get("https://osu.ppy.sh/api/get_beatmaps")
pub(crate) async fn build(self, client: &Client) -> Result<Response> {
Ok(client
.build_request("https://osu.ppy.sh/api/get_beatmaps")
.await?
.query(&self.kind.to_query())
.query(&self.since.map(|v| ("since", v)).to_query())
.query(&self.mode.to_query())
.send()
.await?)
}
}
@ -144,9 +150,10 @@ pub mod builders {
self
}
pub(crate) fn build(&self, client: &Client) -> RequestBuilder {
client
.get("https://osu.ppy.sh/api/get_user")
pub(crate) async fn build(&self, client: &Client) -> Result<Response> {
Ok(client
.build_request("https://osu.ppy.sh/api/get_user")
.await?
.query(&self.user.to_query())
.query(&self.mode.to_query())
.query(
@ -155,6 +162,8 @@ pub mod builders {
.map(|v| ("event_days", v.to_string()))
.to_query(),
)
.send()
.await?)
}
}
@ -197,14 +206,17 @@ pub mod builders {
self
}
pub(crate) fn build(&self, client: &Client) -> RequestBuilder {
client
.get("https://osu.ppy.sh/api/get_scores")
pub(crate) async fn build(&self, client: &Client) -> Result<Response> {
Ok(client
.build_request("https://osu.ppy.sh/api/get_scores")
.await?
.query(&[("b", self.beatmap_id)])
.query(&self.user.to_query())
.query(&self.mode.to_query())
.query(&self.mods.to_query())
.query(&self.limit.map(|v| ("limit", v.to_string())).to_query())
.send()
.await?)
}
}
@ -240,15 +252,18 @@ pub mod builders {
self
}
pub(crate) fn build(&self, client: &Client) -> RequestBuilder {
client
.get(match self.score_type {
pub(crate) async fn build(&self, client: &Client) -> Result<Response> {
Ok(client
.build_request(match self.score_type {
UserScoreType::Best => "https://osu.ppy.sh/api/get_user_best",
UserScoreType::Recent => "https://osu.ppy.sh/api/get_user_recent",
})
.await?
.query(&self.user.to_query())
.query(&self.mode.to_query())
.query(&self.limit.map(|v| ("limit", v.to_string())).to_query())
.send()
.await?)
}
}
}

13
youmubot-prelude/Cargo.toml
View file

@ -7,9 +7,16 @@ edition = "2018"
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
[dependencies]
serenity = "0.8"
anyhow = "1.0"
async-trait = "0.1"
futures-util = "0.3"
tokio = { version = "0.2", features = ["time"] }
youmubot-db = { path = "../youmubot-db" }
crossbeam-channel = "0.4"
reqwest = "0.10"
rayon = "1"
chrono = "0.4"
flume = "0.9"
[dependencies.serenity]
version = "0.9.0-rc.0"
default-features = true
features = ["collector"]

196
youmubot-prelude/src/announcer.rs
View file

@ -1,10 +1,13 @@
use crate::{AppData, GetCloned};
use crossbeam_channel::after;
use rayon::prelude::*;
use crate::{AppData, Result};
use async_trait::async_trait;
use futures_util::{
future::{join_all, ready, FutureExt},
stream::{FuturesUnordered, StreamExt},
};
use serenity::{
framework::standard::{
macros::{command, group},
Args, CommandError as Error, CommandResult,
Args, CommandResult,
},
http::CacheHttp,
model::{
@ -15,11 +18,7 @@ use serenity::{
utils::MessageBuilder,
CacheAndHttp,
};
use std::{
collections::HashMap,
sync::Arc,
thread::{spawn, JoinHandle},
};
use std::{collections::HashMap, sync::Arc};
use youmubot_db::DB;
/// A list of assigned channels for an announcer.
@ -33,30 +32,17 @@ pub(crate) type AnnouncerChannels = DB<HashMap<String, HashMap<GuildId, ChannelI
/// - An AppData, which can be used for interacting with internal databases.
/// - A function "channels", which takes an UserId and returns the list of ChannelIds, which any update related to that user should be
/// sent to.
#[async_trait]
pub trait Announcer: Send {
/// Look for updates and send them to respective channels.
///
/// Errors returned from this function gets ignored and logged down.
fn updates(
async fn updates(
&mut self,
c: Arc<CacheAndHttp>,
d: AppData,
channels: MemberToChannels,
) -> CommandResult;
}
impl<T> Announcer for T
where
T: FnMut(Arc<CacheAndHttp>, AppData, MemberToChannels) -> CommandResult + Send,
{
fn updates(
&mut self,
c: Arc<CacheAndHttp>,
d: AppData,
channels: MemberToChannels,
) -> CommandResult {
self(c, d, channels)
}
) -> Result<()>;
}
/// A simple struct that allows looking up the relevant channels to an user.
@ -64,18 +50,24 @@ pub struct MemberToChannels(Vec<(GuildId, ChannelId)>);
impl MemberToChannels {
/// Gets the channel list of an user related to that channel.
pub fn channels_of(
pub async fn channels_of(
&self,
http: impl CacheHttp + Clone + Sync,
u: impl Into<UserId>,
) -> Vec<ChannelId> {
let u = u.into();
let u: UserId = u.into();
self.0
.par_iter()
.filter_map(|(guild, channel)| {
guild.member(http.clone(), u).ok().map(|_| channel.clone())
.clone()
.into_iter()
.map(|(guild, channel): (GuildId, ChannelId)| {
guild
.member(http.clone(), u)
.map(move |v| v.ok().map(|_| channel.clone()))
})
.collect::<Vec<_>>()
.collect::<FuturesUnordered<_>>()
.filter_map(|v| ready(v))
.collect()
.await
}
}
@ -85,7 +77,7 @@ impl MemberToChannels {
pub struct AnnouncerHandler {
cache_http: Arc<CacheAndHttp>,
data: AppData,
announcers: HashMap<&'static str, Box<dyn Announcer>>,
announcers: HashMap<&'static str, RwLock<Box<dyn Announcer + Send + Sync>>>,
}
// Querying for the AnnouncerHandler in the internal data returns a vec of keys.
@ -107,8 +99,15 @@ impl AnnouncerHandler {
/// Insert a new announcer into the handler.
///
/// The handler must take an unique key. If a duplicate is found, this method panics.
pub fn add(&mut self, key: &'static str, announcer: impl Announcer + 'static) -> &mut Self {
if let Some(_) = self.announcers.insert(key, Box::new(announcer)) {
pub fn add(
&mut self,
key: &'static str,
announcer: impl Announcer + Send + Sync + 'static,
) -> &mut Self {
if let Some(_) = self
.announcers
.insert(key, RwLock::new(Box::new(announcer)))
{
panic!(
"Announcer keys must be unique: another announcer with key `{}` was found",
key
@ -122,9 +121,8 @@ impl AnnouncerHandler {
/// Execution-related.
impl AnnouncerHandler {
/// Collect the list of guilds and their respective channels, by the key of the announcer.
fn get_guilds(&self, key: &'static str) -> Result<Vec<(GuildId, ChannelId)>, Error> {
let d = &self.data;
let data = AnnouncerChannels::open(&*d.read())
async fn get_guilds(data: &AppData, key: &'static str) -> Result<Vec<(GuildId, ChannelId)>> {
let data = AnnouncerChannels::open(&*data.read().await)
.borrow()?
.get(key)
.map(|m| m.iter().map(|(a, b)| (*a, *b)).collect())
@ -133,48 +131,55 @@ impl AnnouncerHandler {
}
/// Run the announcing sequence on a certain announcer.
fn announce(&mut self, key: &'static str) -> CommandResult {
let guilds: Vec<_> = self.get_guilds(key)?;
let channels = MemberToChannels(guilds);
let cache_http = self.cache_http.clone();
let data = self.data.clone();
let announcer = self
.announcers
.get_mut(&key)
.expect("Key is from announcers");
announcer.updates(cache_http, data, channels)?;
Ok(())
async fn announce(
data: AppData,
cache_http: Arc<CacheAndHttp>,
key: &'static str,
announcer: &'_ RwLock<Box<dyn Announcer + Send + Sync>>,
) -> Result<()> {
let channels = MemberToChannels(Self::get_guilds(&data, key).await?);
announcer
.write()
.await
.updates(cache_http, data, channels)
.await
}
/// Start the AnnouncerHandler, moving it into another thread.
/// Start the AnnouncerHandler, looping forever.
///
/// It will run all the announcers in sequence every *cooldown* seconds.
pub fn scan(mut self, cooldown: std::time::Duration) -> JoinHandle<()> {
pub async fn scan(self, cooldown: std::time::Duration) -> () {
// First we store all the keys inside the database.
let keys = self.announcers.keys().cloned().collect::<Vec<_>>();
self.data.write().insert::<Self>(keys.clone());
spawn(move || loop {
self.data.write().await.insert::<Self>(keys.clone());
loop {
eprintln!("{}: announcer started scanning", chrono::Utc::now());
let after_timer = after(cooldown);
for key in &keys {
// let after_timer = after(cooldown);
let after = tokio::time::delay_for(cooldown);
join_all(self.announcers.iter().map(|(key, announcer)| {
eprintln!(" - scanning key `{}`", key);
if let Err(e) = self.announce(key) {
dbg!(e);
}
}
Self::announce(self.data.clone(), self.cache_http.clone(), *key, announcer).map(
move |v| {
if let Err(e) = v {
eprintln!(" - key `{}`: {:?}", *key, e)
}
},
)
}))
.await;
eprintln!("{}: announcer finished scanning", chrono::Utc::now());
after_timer.recv().ok();
})
after.await;
}
}
}
/// Gets the announcer of the given guild.
pub fn announcer_of(
pub async fn announcer_of(
ctx: &Context,
key: &'static str,
guild: GuildId,
) -> Result<Option<ChannelId>, Error> {
Ok(AnnouncerChannels::open(&*ctx.data.read())
) -> Result<Option<ChannelId>> {
Ok(AnnouncerChannels::open(&*ctx.data.read().await)
.borrow()?
.get(key)
.and_then(|channels| channels.get(&guild).cloned()))
@ -184,20 +189,20 @@ pub fn announcer_of(
#[description = "List the registered announcers of this server"]
#[num_args(0)]
#[only_in(guilds)]
pub fn list_announcers(ctx: &mut Context, m: &Message, _: Args) -> CommandResult {
pub async fn list_announcers(ctx: &Context, m: &Message, _: Args) -> CommandResult {
let guild_id = m.guild_id.unwrap();
let announcers = AnnouncerChannels::open(&*ctx.data.read());
let announcers = announcers.borrow()?;
let channels = ctx
.data
.get_cloned::<AnnouncerHandler>()
.into_iter()
.filter_map(|key| {
announcers
.get(key)
.and_then(|channels| channels.get(&guild_id))
.map(|&ch| (key, ch))
let data = &*ctx.data.read().await;
let announcers = AnnouncerChannels::open(data);
let channels = data.get::<AnnouncerHandler>().unwrap();
let channels = channels
.iter()
.filter_map(|&key| {
announcers.borrow().ok().and_then(|announcers| {
announcers
.get(key)
.and_then(|channels| channels.get(&guild_id))
.map(|&ch| (key, ch))
})
})
.map(|(key, ch)| format!(" - `{}`: activated on channel {}", key, ch.mention()))
.collect::<Vec<_>>();
@ -208,7 +213,8 @@ pub fn list_announcers(ctx: &mut Context, m: &Message, _: Args) -> CommandResult
"Activated announcers on this server:\n{}",
channels.join("\n")
),
)?;
)
.await?;
Ok(())
}
@ -219,23 +225,24 @@ pub fn list_announcers(ctx: &mut Context, m: &Message, _: Args) -> CommandResult
#[required_permissions(MANAGE_CHANNELS)]
#[only_in(guilds)]
#[num_args(1)]
pub fn register_announcer(ctx: &mut Context, m: &Message, mut args: Args) -> CommandResult {
pub async fn register_announcer(ctx: &Context, m: &Message, mut args: Args) -> CommandResult {
let data = ctx.data.read().await;
let key = args.single::<String>()?;
let keys = ctx.data.get_cloned::<AnnouncerHandler>();
if !keys.contains(&key.as_str()) {
let keys = data.get::<AnnouncerHandler>().unwrap();
if !keys.contains(&&key[..]) {
m.reply(
&ctx,
format!(
"Key not found. Available announcer keys are: `{}`",
keys.join(", ")
),
)?;
)
.await?;
return Ok(());
}
let guild = m.guild(&ctx).expect("Guild-only command");
let guild = guild.read();
let channel = m.channel_id.to_channel(&ctx)?;
AnnouncerChannels::open(&*ctx.data.read())
let guild = m.guild(&ctx).await.expect("Guild-only command");
let channel = m.channel_id.to_channel(&ctx).await?;
AnnouncerChannels::open(&*data)
.borrow_mut()?
.entry(key.clone())
.or_default()
@ -250,7 +257,8 @@ pub fn register_announcer(ctx: &mut Context, m: &Message, mut args: Args) -> Com
.push(" on channel ")
.push_bold_safe(channel)
.build(),
)?;
)
.await?;
Ok(())
}
@ -260,9 +268,10 @@ pub fn register_announcer(ctx: &mut Context, m: &Message, mut args: Args) -> Com
#[required_permissions(MANAGE_CHANNELS)]
#[only_in(guilds)]
#[num_args(1)]
pub fn remove_announcer(ctx: &mut Context, m: &Message, mut args: Args) -> CommandResult {
pub async fn remove_announcer(ctx: &Context, m: &Message, mut args: Args) -> CommandResult {
let data = ctx.data.read().await;
let key = args.single::<String>()?;
let keys = ctx.data.get_cloned::<AnnouncerHandler>();
let keys = data.get::<AnnouncerHandler>().unwrap();
if !keys.contains(&key.as_str()) {
m.reply(
&ctx,
@ -270,12 +279,12 @@ pub fn remove_announcer(ctx: &mut Context, m: &Message, mut args: Args) -> Comma
"Key not found. Available announcer keys are: `{}`",
keys.join(", ")
),
)?;
)
.await?;
return Ok(());
}
let guild = m.guild(&ctx).expect("Guild-only command");
let guild = guild.read();
AnnouncerChannels::open(&*ctx.data.read())
let guild = m.guild(&ctx).await.expect("Guild-only command");
AnnouncerChannels::open(&*data)
.borrow_mut()?
.entry(key.clone())
.and_modify(|m| {
@ -289,7 +298,8 @@ pub fn remove_announcer(ctx: &mut Context, m: &Message, mut args: Args) -> Comma
.push(" has been de-activated for server ")
.push_bold_safe(&guild.name)
.build(),
)?;
)
.await?;
Ok(())
}

17
youmubot-prelude/src/args.rs
View file

@ -2,14 +2,17 @@ pub use duration::Duration;
pub use username_arg::UsernameArg;
mod duration {
use crate::{Error, Result};
use std::fmt;
use std::time::Duration as StdDuration;
use String as Error;
// Parse a single duration unit
fn parse_duration_string(s: &str) -> Result<StdDuration, Error> {
const INVALID_DURATION: &str = "Not a valid duration";
/// Parse a single duration unit
fn parse_duration_string(s: &str) -> Result<StdDuration> {
// We reject the empty case
if s == "" {
return Err(Error::from("empty strings are not valid durations"));
return Err(Error::msg("empty strings are not valid durations"));
}
struct ParseStep {
current_value: Option<u64>,
@ -26,7 +29,7 @@ mod duration {
current_value: Some(v.unwrap_or(0) * 10 + ((item as u64) - ('0' as u64))),
..s
}),
(_, None) => Err(Error::from("Not a valid duration")),
(_, None) => Err(Error::msg(INVALID_DURATION)),
(item, Some(v)) => Ok(ParseStep {
current_value: None,
current_duration: s.current_duration
@ -36,7 +39,7 @@ mod duration {
'h' => StdDuration::from_secs(60 * 60),
'd' => StdDuration::from_secs(60 * 60 * 24),
'w' => StdDuration::from_secs(60 * 60 * 24 * 7),
_ => return Err(Error::from("Not a valid duration")),
_ => return Err(Error::msg(INVALID_DURATION)),
} * (v as u32),
}),
},
@ -44,7 +47,7 @@ mod duration {
.and_then(|v| match v.current_value {
// All values should be consumed
None => Ok(v),
_ => Err(Error::from("Not a valid duration")),
_ => Err(Error::msg(INVALID_DURATION)),
})
.map(|v| v.current_duration)
}

24
youmubot-prelude/src/hook.rs Normal file
View file

@ -0,0 +1,24 @@
use crate::{async_trait, future, Context, Result};
use serenity::model::channel::Message;
/// Hook represents the asynchronous hook that is run on every message.
#[async_trait]
pub trait Hook: Send + Sync {
async fn call(&mut self, ctx: &Context, message: &Message) -> Result<()>;
}
#[async_trait]
impl<T> Hook for T
where
T: for<'a> FnMut(
&'a Context,
&'a Message,
)
-> std::pin::Pin<Box<dyn future::Future<Output = Result<()>> + 'a + Send>>
+ Send
+ Sync,
{
async fn call(&mut self, ctx: &Context, message: &Message) -> Result<()> {
self(ctx, message).await
}
}

58
youmubot-prelude/src/lib.rs
View file

@ -1,54 +1,40 @@
/// Module `prelude` provides a sane set of default imports that can be used inside
/// a Youmubot source file.
pub use serenity::prelude::*;
use std::sync::Arc;
pub mod announcer;
pub mod args;
pub mod hook;
pub mod pagination;
pub mod reaction_watch;
pub mod ratelimit;
pub mod setup;
pub use announcer::{Announcer, AnnouncerHandler};
pub use args::{Duration, UsernameArg};
pub use pagination::Pagination;
pub use reaction_watch::{ReactionHandler, ReactionWatcher};
pub use hook::Hook;
pub use pagination::paginate;
/// Re-exporting async_trait helps with implementing Announcer.
pub use async_trait::async_trait;
/// Re-export the anyhow errors
pub use anyhow::{Error, Result};
/// Re-export useful future and stream utils
pub use futures_util::{future, stream, FutureExt, StreamExt, TryFutureExt, TryStreamExt};
/// Re-export the spawn function
pub use tokio::spawn as spawn_future;
/// The global app data.
pub type AppData = Arc<RwLock<ShareMap>>;
pub type AppData = Arc<RwLock<TypeMap>>;
/// The HTTP client.
pub struct HTTPClient;
impl TypeMapKey for HTTPClient {
type Value = reqwest::blocking::Client;
}
/// The TypeMap trait that allows TypeMaps to quickly get a clonable item.
pub trait GetCloned {
/// Gets an item from the store, cloned.
fn get_cloned<T>(&self) -> T::Value
where
T: TypeMapKey,
T::Value: Clone + Send + Sync;
}
impl GetCloned for ShareMap {
fn get_cloned<T>(&self) -> T::Value
where
T: TypeMapKey,
T::Value: Clone + Send + Sync,
{
self.get::<T>().cloned().expect("Should be there")
}
}
impl GetCloned for AppData {
fn get_cloned<T>(&self) -> T::Value
where
T: TypeMapKey,
T::Value: Clone + Send + Sync,
{
self.read().get::<T>().cloned().expect("Should be there")
}
type Value = reqwest::Client;
}
pub mod prelude_commands {
@ -70,8 +56,8 @@ pub mod prelude_commands {
#[command]
#[description = "pong!"]
fn ping(ctx: &mut Context, m: &Message) -> CommandResult {
m.reply(&ctx, "Pong!")?;
async fn ping(ctx: &Context, m: &Message) -> CommandResult {
m.reply(&ctx, "Pong!").await?;
Ok(())
}
}

230
youmubot-prelude/src/pagination.rs
View file

@ -1,157 +1,111 @@
use crate::{Context, ReactionHandler, ReactionWatcher};
use crate::{Context, Result};
use futures_util::{future::Future, StreamExt};
use serenity::{
builder::EditMessage,
framework::standard::{CommandError, CommandResult},
collector::ReactionAction,
model::{
channel::{Message, Reaction, ReactionType},
channel::{Message, ReactionType},
id::ChannelId,
},
};
use std::convert::TryFrom;
use tokio::time as tokio_time;
const ARROW_RIGHT: &'static str = "➡️";
const ARROW_LEFT: &'static str = "⬅️";
impl ReactionWatcher {
/// Start a pagination.
///
/// Takes a copy of Context (which you can `clone`), a pager (see "Pagination") and a target channel id.
/// Pagination will handle all events on adding/removing an "arrow" emoji (⬅️ and ➡️).
/// This is a blocking call - it will block the thread until duration is over.
pub fn paginate<T: Pagination + Send + 'static>(
&self,
ctx: Context,
channel: ChannelId,
pager: T,
duration: std::time::Duration,
) -> CommandResult {
let handler = PaginationHandler::new(pager, ctx, channel)?;
self.handle_reactions(handler, duration, |_| {});
Ok(())
}
/// A version of `paginate` that compiles for closures.
///
/// A workaround until https://github.com/rust-lang/rust/issues/36582 is solved.
pub fn paginate_fn<T>(
&self,
ctx: Context,
channel: ChannelId,
pager: T,
duration: std::time::Duration,
) -> CommandResult
where
T: for<'a> FnMut(u8, &'a mut EditMessage) -> (&'a mut EditMessage, CommandResult)
+ Send
+ 'static,
{
self.paginate(ctx, channel, pager, duration)
}
#[async_trait::async_trait]
pub trait Paginate {
async fn render(&mut self, page: u8, ctx: &Context, m: &mut Message) -> Result<bool>;
}
/// Pagination allows the bot to display content in multiple pages.
///
/// You need to implement the "render_page" function, which takes a dummy content and
/// embed assigning function.
/// Pagination is automatically implemented for functions with the same signature as `render_page`.
///
/// Pages start at 0.
pub trait Pagination {
/// Render a page.
///
/// This would either create or edit a message, but you should not be worry about it.
fn render_page<'a>(
&mut self,
page: u8,
target: &'a mut EditMessage,
) -> (&'a mut EditMessage, CommandResult);
}
impl<T> Pagination for T
#[async_trait::async_trait]
impl<T> Paginate for T
where
T: for<'a> FnMut(u8, &'a mut EditMessage) -> (&'a mut EditMessage, CommandResult),
T: for<'m> FnMut(
u8,
&'m Context,
&'m mut Message,
) -> std::pin::Pin<Box<dyn Future<Output = Result<bool>> + Send + 'm>>
+ Send,
{
fn render_page<'a>(
&mut self,
page: u8,
target: &'a mut EditMessage,
) -> (&'a mut EditMessage, CommandResult) {
self(page, target)
async fn render(&mut self, page: u8, ctx: &Context, m: &mut Message) -> Result<bool> {
self(page, ctx, m).await
}
}
struct PaginationHandler<T: Pagination> {
pager: T,
message: Message,
// Paginate! with a pager function.
/// If awaited, will block until everything is done.
pub async fn paginate(
mut pager: impl for<'m> FnMut(
u8,
&'m Context,
&'m mut Message,
) -> std::pin::Pin<Box<dyn Future<Output = Result<bool>> + Send + 'm>>
+ Send,
ctx: &Context,
channel: ChannelId,
timeout: std::time::Duration,
) -> Result<()> {
let mut message = channel
.send_message(&ctx, |e| e.content("Youmu is loading the first page..."))
.await?;
// React to the message
message
.react(&ctx, ReactionType::try_from(ARROW_LEFT)?)
.await?;
message
.react(&ctx, ReactionType::try_from(ARROW_RIGHT)?)
.await?;
pager(0, ctx, &mut message).await?;
// Build a reaction collector
let mut reaction_collector = message.await_reactions(&ctx).removed(true).await;
let mut page = 0;
// Loop the handler function.
let res: Result<()> = loop {
match tokio_time::timeout(timeout, reaction_collector.next()).await {
Err(_) => break Ok(()),
Ok(None) => break Ok(()),
Ok(Some(reaction)) => {
page = match handle_reaction(page, &mut pager, ctx, &mut message, &reaction).await {
Ok(v) => v,
Err(e) => break Err(e),
};
}
}
};
message.react(&ctx, '🛑').await?;
res
}
// Handle the reaction and return a new page number.
async fn handle_reaction(
page: u8,
ctx: Context,
}
impl<T: Pagination> PaginationHandler<T> {
pub fn new(pager: T, mut ctx: Context, channel: ChannelId) -> Result<Self, CommandError> {
let message = channel.send_message(&mut ctx, |e| {
e.content("Youmu is loading the first page...")
})?;
// React to the message
message.react(&mut ctx, ARROW_LEFT)?;
message.react(&mut ctx, ARROW_RIGHT)?;
let mut p = Self {
pager,
message: message.clone(),
page: 0,
ctx,
};
p.call_pager()?;
Ok(p)
}
}
impl<T: Pagination> PaginationHandler<T> {
/// Call the pager, log the error (if any).
fn call_pager(&mut self) -> CommandResult {
let mut res: CommandResult = Ok(());
let mut msg = self.message.clone();
msg.edit(self.ctx.http.clone(), |e| {
let (e, r) = self.pager.render_page(self.page, e);
res = r;
e
})?;
self.message = msg;
res
}
}
impl<T: Pagination> Drop for PaginationHandler<T> {
fn drop(&mut self) {
self.message.react(&self.ctx, "🛑").ok();
}
}
impl<T: Pagination> ReactionHandler for PaginationHandler<T> {
fn handle_reaction(&mut self, reaction: &Reaction, _is_add: bool) -> CommandResult {
if reaction.message_id != self.message.id {
return Ok(());
}
match &reaction.emoji {
ReactionType::Unicode(ref s) => match s.as_str() {
ARROW_LEFT if self.page == 0 => return Ok(()),
ARROW_LEFT => {
self.page -= 1;
if let Err(e) = self.call_pager() {
self.page += 1;
return Err(e);
}
}
ARROW_RIGHT => {
self.page += 1;
if let Err(e) = self.call_pager() {
self.page -= 1;
return Err(e);
}
}
_ => (),
},
_ => (),
}
Ok(())
pager: &mut impl Paginate,
ctx: &Context,
message: &mut Message,
reaction: &ReactionAction,
) -> Result<u8> {
let reaction = match reaction {
ReactionAction::Added(v) | ReactionAction::Removed(v) => v,
};
match &reaction.emoji {
ReactionType::Unicode(ref s) => match s.as_str() {
ARROW_LEFT if page == 0 => Ok(page),
ARROW_LEFT => Ok(if pager.render(page - 1, ctx, message).await? {
page - 1
} else {
page
}),
ARROW_RIGHT => Ok(if pager.render(page + 1, ctx, message).await? {
page + 1
} else {
page
}),
_ => Ok(page),
},
_ => Ok(page),
}
}

67
youmubot-prelude/src/ratelimit.rs Normal file
View file

@ -0,0 +1,67 @@
/// Provides a simple ratelimit lock (that only works in tokio)
// use tokio::time::
use std::time::Duration;
use crate::Result;
use flume::{bounded as channel, Receiver, Sender};
use std::ops::Deref;
/// Holds the underlying `T` in a rate-limited way.
pub struct Ratelimit<T> {
inner: T,
recv: Receiver<()>,
send: Sender<()>,
wait_time: Duration,
}
struct RatelimitGuard<'a, T> {
inner: &'a T,
send: &'a Sender<()>,
wait_time: &'a Duration,
}
impl<T> Ratelimit<T> {
/// Create a new ratelimit with at most `count` uses in `wait_time`.
pub fn new(inner: T, count: usize, wait_time: Duration) -> Self {
let (send, recv) = channel(count);
(0..count).for_each(|_| {
send.send(()).ok();
});
Self {
inner,
send,
recv,
wait_time,
}
}
/// Borrow the inner `T`. You can only hol this reference `count` times in `wait_time`.
/// The clock counts from the moment the ref is dropped.
pub async fn borrow<'a>(&'a self) -> Result<impl Deref<Target = T> + 'a> {
self.recv.recv_async().await?;
Ok(RatelimitGuard {
inner: &self.inner,
send: &self.send,
wait_time: &self.wait_time,
})
}
}
impl<'a, T> Deref for RatelimitGuard<'a, T> {
type Target = T;
fn deref(&self) -> &Self::Target {
self.inner
}
}
impl<'a, T> Drop for RatelimitGuard<'a, T> {
fn drop(&mut self) {
let send = self.send.clone();
let wait_time = self.wait_time.clone();
tokio::spawn(async move {
tokio::time::delay_for(wait_time).await;
send.send_async(()).await.ok();
});
}
}

105
youmubot-prelude/src/reaction_watch.rs
View file

@ -1,105 +0,0 @@
use crossbeam_channel::{after, bounded, select, Sender};
use serenity::{framework::standard::CommandResult, model::channel::Reaction, prelude::*};
use std::sync::{Arc, Mutex};
/// Handles a reaction.
///
/// Every handler needs an expire time too.
pub trait ReactionHandler {
/// Handle a reaction. This is fired on EVERY reaction.
/// You do the filtering yourself.
///
/// If `is_added` is false, the reaction was removed instead of added.
fn handle_reaction(&mut self, reaction: &Reaction, is_added: bool) -> CommandResult;
}
impl<T> ReactionHandler for T
where
T: FnMut(&Reaction, bool) -> CommandResult,
{
fn handle_reaction(&mut self, reaction: &Reaction, is_added: bool) -> CommandResult {
self(reaction, is_added)
}
}
/// The store for a set of dynamic reaction handlers.
#[derive(Debug, Clone)]
pub struct ReactionWatcher {
channels: Arc<Mutex<Vec<Sender<(Arc<Reaction>, bool)>>>>,
}
impl TypeMapKey for ReactionWatcher {
type Value = ReactionWatcher;
}
impl ReactionWatcher {
/// Create a new ReactionWatcher.
pub fn new() -> Self {
Self {
channels: Arc::new(Mutex::new(vec![])),
}
}
/// Send a reaction.
/// If `is_added` is false, the reaction was removed.
pub fn send(&self, r: Reaction, is_added: bool) {
let r = Arc::new(r);
self.channels
.lock()
.expect("Poisoned!")
.retain(|e| e.send((r.clone(), is_added)).is_ok());
}
/// React! to a series of reaction
///
/// The reactions stop after `duration` of idle.
pub fn handle_reactions<H: ReactionHandler + Send + 'static>(
&self,
mut h: H,
duration: std::time::Duration,
callback: impl FnOnce(H) -> () + Send + 'static,
) {
let (send, reactions) = bounded(0);
{
self.channels.lock().expect("Poisoned!").push(send);
}
std::thread::spawn(move || {
loop {
let timeout = after(duration);
let r = select! {
recv(reactions) -> r => { let (r, is_added) = r.unwrap(); h.handle_reaction(&*r, is_added) },
recv(timeout) -> _ => break,
};
if let Err(v) = r {
dbg!(v);
}
}
callback(h)
});
}
/// React! to a series of reaction
///
/// The handler will stop after `duration` no matter what.
pub fn handle_reactions_timed<H: ReactionHandler + Send + 'static>(
&self,
mut h: H,
duration: std::time::Duration,
callback: impl FnOnce(H) -> () + Send + 'static,
) {
let (send, reactions) = bounded(0);
{
self.channels.lock().expect("Poisoned!").push(send);
}
std::thread::spawn(move || {
let timeout = after(duration);
loop {
let r = select! {
recv(reactions) -> r => { let (r, is_added) = r.unwrap(); h.handle_reaction(&*r, is_added) },
recv(timeout) -> _ => break,
};
if let Err(v) = r {
dbg!(v);
}
}
callback(h);
});
}
}

9
youmubot-prelude/src/setup.rs
View file

@ -1,17 +1,14 @@
use serenity::{framework::standard::StandardFramework, prelude::*};
use serenity::prelude::*;
use std::path::Path;
/// Set up the prelude libraries.
///
/// Panics on failure: Youmubot should *NOT* attempt to continue when this function fails.
pub fn setup_prelude(db_path: &Path, data: &mut ShareMap, _: &mut StandardFramework) {
pub fn setup_prelude(db_path: &Path, data: &mut TypeMap) {
// Setup the announcer DB.
crate::announcer::AnnouncerChannels::insert_into(data, db_path.join("announcers.yaml"))
.expect("Announcers DB set up");
// Set up the HTTP client.
data.insert::<crate::HTTPClient>(reqwest::blocking::Client::new());
// Set up the reaction watcher.
data.insert::<crate::ReactionWatcher>(crate::ReactionWatcher::new());
data.insert::<crate::HTTPClient>(reqwest::Client::new());
}

4
youmubot/Cargo.toml
View file

@ -12,8 +12,10 @@ osu = ["youmubot-osu"]
codeforces = ["youmubot-cf"]
[dependencies]
serenity = "0.8"
serenity = "0.9.0-rc.0"
tokio = "0.2"
dotenv = "0.15"
env_logger = "0.7"
youmubot-db = { path = "../youmubot-db" }
youmubot-prelude = { path = "../youmubot-prelude" }
youmubot-core = { path = "../youmubot-core" }

246
youmubot/src/main.rs
View file

@ -1,9 +1,10 @@
use dotenv;
use dotenv::var;
use serenity::{
framework::standard::{DispatchError, StandardFramework},
client::bridge::gateway::GatewayIntents,
framework::standard::{macros::hook, CommandResult, DispatchError, StandardFramework},
model::{
channel::{Channel, Message, Reaction},
channel::{Channel, Message},
gateway,
id::{ChannelId, GuildId, UserId},
permissions::Permissions,
@ -12,51 +13,59 @@ use serenity::{
use youmubot_prelude::*;
struct Handler {
hooks: Vec<fn(&mut Context, &Message) -> ()>,
hooks: Vec<RwLock<Box<dyn Hook>>>,
}
impl Handler {
fn new() -> Handler {
Handler { hooks: vec![] }
}
fn push_hook<T: Hook + 'static>(&mut self, f: T) {
self.hooks.push(RwLock::new(Box::new(f)));
}
}
#[async_trait]
impl EventHandler for Handler {
fn ready(&self, _: Context, ready: gateway::Ready) {
async fn ready(&self, _: Context, ready: gateway::Ready) {
println!("{} is connected!", ready.user.name);
}
fn message(&self, mut ctx: Context, message: Message) {
self.hooks.iter().for_each(|f| f(&mut ctx, &message));
}
fn reaction_add(&self, ctx: Context, reaction: Reaction) {
ctx.data
.get_cloned::<ReactionWatcher>()
.send(reaction, true);
}
fn reaction_remove(&self, ctx: Context, reaction: Reaction) {
ctx.data
.get_cloned::<ReactionWatcher>()
.send(reaction, false);
async fn message(&self, ctx: Context, message: Message) {
self.hooks
.iter()
.map(|hook| {
let ctx = ctx.clone();
let message = message.clone();
hook.write()
.then(|mut h| async move { h.call(&ctx, &message).await })
})
.collect::<stream::FuturesUnordered<_>>()
.for_each(|v| async move {
if let Err(e) = v {
eprintln!("{}", e)
}
})
.await;
}
}
/// Returns whether the user has "MANAGE_MESSAGES" permission in the channel.
fn is_channel_mod(ctx: &mut Context, _: Option<GuildId>, ch: ChannelId, u: UserId) -> bool {
match ch.to_channel(&ctx) {
Ok(Channel::Guild(gc)) => {
let gc = gc.read();
gc.permissions_for_user(&ctx, u)
.map(|perms| perms.contains(Permissions::MANAGE_MESSAGES))
.unwrap_or(false)
}
async fn is_channel_mod(ctx: &Context, _: Option<GuildId>, ch: ChannelId, u: UserId) -> bool {
match ch.to_channel(&ctx).await {
Ok(Channel::Guild(gc)) => gc
.permissions_for_user(&ctx, u)
.await
.map(|perms| perms.contains(Permissions::MANAGE_MESSAGES))
.unwrap_or(false),
_ => false,
}
}
fn main() {
#[tokio::main]
async fn main() {
env_logger::init();
// Setup dotenv
if let Ok(path) = dotenv::dotenv() {
println!("Loaded dotenv from {:?}", path);
@ -65,34 +74,48 @@ fn main() {
let mut handler = Handler::new();
// Set up hooks
#[cfg(feature = "osu")]
handler.hooks.push(youmubot_osu::discord::hook);
handler.push_hook(youmubot_osu::discord::hook);
#[cfg(feature = "codeforces")]
handler.hooks.push(youmubot_cf::codeforces_info_hook);
handler.push_hook(youmubot_cf::InfoHook);
// Collect the token
let token = var("TOKEN").expect("Please set TOKEN as the Discord Bot's token to be used.");
// Set up base framework
let fw = setup_framework(&token[..]).await;
// Sets up a client
let mut client = {
// Collect the token
let token = var("TOKEN").expect("Please set TOKEN as the Discord Bot's token to be used.");
// Attempt to connect and set up a framework
Client::new(token, handler).expect("Cannot connect")
Client::new(token)
.framework(fw)
.event_handler(handler)
.intents(
GatewayIntents::GUILDS
| GatewayIntents::GUILD_BANS
| GatewayIntents::GUILD_MESSAGES
| GatewayIntents::GUILD_MESSAGE_REACTIONS
| GatewayIntents::GUILD_PRESENCES
| GatewayIntents::GUILD_MEMBERS
| GatewayIntents::DIRECT_MESSAGES
| GatewayIntents::DIRECT_MESSAGE_REACTIONS,
)
.await
.unwrap()
};
// Set up base framework
let mut fw = setup_framework(&client);
// Set up announcer handler
let mut announcers = AnnouncerHandler::new(&client);
// Setup each package starting from the prelude.
{
let mut data = client.data.write();
let mut data = client.data.write().await;
let db_path = var("DBPATH")
.map(|v| std::path::PathBuf::from(v))
.unwrap_or_else(|e| {
println!("No DBPATH set up ({:?}), using `/data`", e);
std::path::PathBuf::from("data")
});
youmubot_prelude::setup::setup_prelude(&db_path, &mut data, &mut fw);
youmubot_prelude::setup::setup_prelude(&db_path, &mut data);
// Setup core
#[cfg(feature = "core")]
youmubot_core::setup(&db_path, &client, &mut data).expect("Setup db should succeed");
@ -102,7 +125,7 @@ fn main() {
.expect("osu! is initialized");
// codeforces
#[cfg(feature = "codeforces")]
youmubot_cf::setup(&db_path, &mut data, &mut announcers);
youmubot_cf::setup(&db_path, &mut data, &mut announcers).await;
}
#[cfg(feature = "core")]
@ -112,11 +135,10 @@ fn main() {
#[cfg(feature = "codeforces")]
println!("codeforces enabled.");
client.with_framework(fw);
announcers.scan(std::time::Duration::from_secs(120));
tokio::spawn(announcers.scan(std::time::Duration::from_secs(120)));
println!("Starting...");
if let Err(v) = client.start() {
if let Err(v) = client.start().await {
panic!(v)
}
@ -124,71 +146,43 @@ fn main() {
}
// Sets up a framework for a client
fn setup_framework(client: &Client) -> StandardFramework {
async fn setup_framework(token: &str) -> StandardFramework {
let http = serenity::http::Http::new_with_token(token);
// Collect owners
let owner = client
.cache_and_http
.http
let owner = http
.get_current_application_info()
.await
.expect("Should be able to get app info")
.owner;
let fw = StandardFramework::new()
.configure(|c| {
c.with_whitespace(false)
.prefix(&var("PREFIX").unwrap_or("y!".to_owned()))
.delimiters(vec![" / ", "/ ", " /", "/"])
.owners([owner.id].iter().cloned().collect())
})
.help(&youmubot_core::HELP)
.before(|_, msg, command_name| {
println!(
"Got command '{}' by user '{}'",
command_name, msg.author.name
);
true
})
.after(|ctx, msg, command_name, error| match error {
Ok(()) => println!("Processed command '{}'", command_name),
Err(why) => {
let reply = format!("Command '{}' returned error {:?}", command_name, why);
if let Err(_) = msg.reply(&ctx, &reply) {}
println!("{}", reply)
}
})
.on_dispatch_error(|ctx, msg, error| {
msg.reply(
&ctx,
&match error {
DispatchError::Ratelimited(seconds) => format!(
"⏳ You are being rate-limited! Try this again in **{} seconds**.",
seconds
),
DispatchError::NotEnoughArguments { min, given } => format!("😕 The command needs at least **{}** arguments, I only got **{}**!\nDid you know command arguments are separated with a slash (`/`)?", min, given),
DispatchError::TooManyArguments { max, given } => format!("😕 I can only handle at most **{}** arguments, but I got **{}**!", max, given),
DispatchError::OnlyForGuilds => format!("🔇 This command cannot be used in DMs."),
_ => return,
},
)
.unwrap(); // Invoke
})
// Set a function that's called whenever an attempted command-call's
// command could not be found.
.unrecognised_command(|_, _, unknown_command_name| {
println!("Could not find command named '{}'", unknown_command_name);
})
// Set a function that's called whenever a message is not a command.
.normal_message(|_, _| {
// println!("Message is not a command '{}'", message.content);
})
.bucket("voting", |c| {
c.check(|ctx, g, ch, u| !is_channel_mod(ctx, g, ch, u)).delay(120 /* 2 minutes */).time_span(120).limit(1)
})
.bucket("images", |c| c.time_span(60).limit(2))
.bucket("community", |c| {
c.check(|ctx, g, ch, u| !is_channel_mod(ctx, g, ch, u)).delay(30).time_span(30).limit(1)
})
.group(&prelude_commands::PRELUDE_GROUP);
let fw = StandardFramework::new()
.configure(|c| {
c.with_whitespace(false)
.prefix(&var("PREFIX").unwrap_or("y!".to_owned()))
.delimiters(vec![" / ", "/ ", " /", "/"])
.owners([owner.id].iter().cloned().collect())
})
.help(&youmubot_core::HELP)
.before(before_hook)
.after(after_hook)
.on_dispatch_error(on_dispatch_error)
.bucket("voting", |c| {
c.check(|ctx, g, ch, u| Box::pin(async move { !is_channel_mod(ctx, g, ch, u).await }))
.delay(120 /* 2 minutes */)
.time_span(120)
.limit(1)
})
.await
.bucket("images", |c| c.time_span(60).limit(2))
.await
.bucket("community", |c| {
c.check(|ctx, g, ch, u| Box::pin(async move { !is_channel_mod(ctx, g, ch, u).await }))
.delay(30)
.time_span(30)
.limit(1)
})
.await
.group(&prelude_commands::PRELUDE_GROUP);
// groups here
#[cfg(feature = "core")]
let fw = fw
@ -201,3 +195,53 @@ fn setup_framework(client: &Client) -> StandardFramework {
let fw = fw.group(&youmubot_cf::CODEFORCES_GROUP);
fw
}
// Hooks!
#[hook]
async fn before_hook(_: &Context, msg: &Message, command_name: &str) -> bool {
println!(
"Got command '{}' by user '{}'",
command_name, msg.author.name
);
true
}
#[hook]
async fn after_hook(ctx: &Context, msg: &Message, command_name: &str, error: CommandResult) {
match error {
Ok(()) => println!("Processed command '{}'", command_name),
Err(why) => {
let reply = format!("Command '{}' returned error {:?}", command_name, why);
msg.reply(&ctx, &reply).await.ok();
println!("{}", reply)
}
}
}
#[hook]
async fn on_dispatch_error(ctx: &Context, msg: &Message, error: DispatchError) {
msg.reply(
&ctx,
&match error {
DispatchError::Ratelimited(seconds) => format!(
"⏳ You are being rate-limited! Try this again in **{}**.",
youmubot_prelude::Duration(seconds),
),
DispatchError::NotEnoughArguments { min, given } => {
format!(
"😕 The command needs at least **{}** arguments, I only got **{}**!",
min, given
) + "\nDid you know command arguments are separated with a slash (`/`)?"
}
DispatchError::TooManyArguments { max, given } => format!(
"😕 I can only handle at most **{}** arguments, but I got **{}**!",
max, given
),
DispatchError::OnlyForGuilds => format!("🔇 This command cannot be used in DMs."),
_ => return,
},
)
.await
.ok(); // Invoke
}