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

Codeforces: asyncify lib

Browse source
This commit is contained in:
Natsu Kagami 2020-09-07 23:21:17 -04:00
parent 4a2a4ed688
commit 80afe79904
Signed by: nki
GPG key ID: 73376E117CD20735
3 changed files with 235 additions and 187 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -9,8 +9,6 @@ use serenity::{http::CacheHttp, model::id::UserId, CacheAndHttp};
use std::sync::Arc; use std::sync::Arc;
use youmubot_prelude::*; use youmubot_prelude::*;
type Reqwest = <HTTPClient as TypeMapKey>::Value;
/// Updates the rating and rating changes of the users. /// Updates the rating and rating changes of the users.
pub struct Announcer; pub struct Announcer;

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

@ -1,9 +1,7 @@
use crate::CFClient;
use chrono::{TimeZone, Utc}; use chrono::{TimeZone, Utc};
use codeforces::{Client, Contest, Problem}; use codeforces::{Client, Contest, Problem};
use dashmap::DashMap as HashMap; use dashmap::DashMap as HashMap;
use lazy_static::lazy_static; use lazy_static::lazy_static;
use rayon::{iter::Either, prelude::*};
use regex::{Captures, Regex}; use regex::{Captures, Regex};
use serenity::{ use serenity::{
builder::CreateEmbed, framework::standard::CommandError, model::channel::Message, builder::CreateEmbed, framework::standard::CommandError, model::channel::Message,
@ -104,7 +102,6 @@ impl Hook for InfoHook {
return Ok(()); return Ok(());
} }
let data = ctx.data.read().await; let data = ctx.data.read().await;
let http = data.get::<CFClient>().unwrap();
let contest_cache = data.get::<ContestCache>().unwrap(); let contest_cache = data.get::<ContestCache>().unwrap();
let matches = parse(&m.content[..], contest_cache) let matches = parse(&m.content[..], contest_cache)
.collect::<Vec<_>>() .collect::<Vec<_>>()
@ -138,11 +135,24 @@ fn print_info_message<'a>(
info: &[(ContestOrProblem, &str)], info: &[(ContestOrProblem, &str)],
e: &'a mut CreateEmbed, e: &'a mut CreateEmbed,
) -> &'a mut CreateEmbed { ) -> &'a mut CreateEmbed {
let (mut problems, contests): (Vec<_>, Vec<_>) = let (problems, contests): (Vec<_>, Vec<_>) = info.iter().partition(|(v, _)| match v {
info.par_iter().partition_map(|(v, l)| match v { ContestOrProblem::Problem(_) => true,
ContestOrProblem::Problem(p) => Either::Left((p, l)), ContestOrProblem::Contest(_, _) => false,
ContestOrProblem::Contest(c, p) => Either::Right((c, p, l)), });
}); 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))); problems.sort_by(|(a, _), (b, _)| a.rating.unwrap_or(1500).cmp(&b.rating.unwrap_or(1500)));
let mut m = MessageBuilder::new(); let mut m = MessageBuilder::new();
if !problems.is_empty() { if !problems.is_empty() {

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

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