From 35471a4a0be367d183f4b853d68f8c06685f9e98 Mon Sep 17 00:00:00 2001 From: Maybe Waffle Date: Mon, 13 Feb 2023 14:11:06 +0400 Subject: [PATCH] Add `mentioned_users` to update types where it makes sense --- crates/teloxide-core/src/lib.rs | 1 + .../teloxide-core/src/types/callback_query.rs | 14 +++++ crates/teloxide-core/src/types/chat.rs | 19 ++++++- .../src/types/chat_join_request.rs | 12 ++++ .../src/types/chat_member_updated.rs | 18 ++++++ crates/teloxide-core/src/types/game.rs | 16 +++++- crates/teloxide-core/src/types/message.rs | 36 ++++++++++++ crates/teloxide-core/src/types/poll.rs | 16 +++++- crates/teloxide-core/src/types/update.rs | 42 ++++++++++++++ crates/teloxide-core/src/util.rs | 56 +++++++++++++++++++ 10 files changed, 227 insertions(+), 3 deletions(-) create mode 100644 crates/teloxide-core/src/util.rs diff --git a/crates/teloxide-core/src/lib.rs b/crates/teloxide-core/src/lib.rs index bfc94ddf..6127ad54 100644 --- a/crates/teloxide-core/src/lib.rs +++ b/crates/teloxide-core/src/lib.rs @@ -117,6 +117,7 @@ mod bot; // implementation details mod serde_multipart; +mod util; #[cfg(test)] mod codegen; diff --git a/crates/teloxide-core/src/types/callback_query.rs b/crates/teloxide-core/src/types/callback_query.rs index 5cfa0fa8..2d1f5362 100644 --- a/crates/teloxide-core/src/types/callback_query.rs +++ b/crates/teloxide-core/src/types/callback_query.rs @@ -49,6 +49,20 @@ pub struct CallbackQuery { pub game_short_name: Option, } +impl CallbackQuery { + /// Returns all users that are "contained" in this `CallbackQuery` + /// structure. + /// + /// This might be useful to track information about users. + /// Note that this function can return duplicate users. + pub fn mentioned_users(&self) -> impl Iterator { + use crate::util::flatten; + use std::iter::once; + + once(&self.from).chain(flatten(self.message.as_ref().map(Message::mentioned_users))) + } +} + #[cfg(test)] mod tests { use crate::types::UserId; diff --git a/crates/teloxide-core/src/types/chat.rs b/crates/teloxide-core/src/types/chat.rs index 38da326b..8dbb4d8d 100644 --- a/crates/teloxide-core/src/types/chat.rs +++ b/crates/teloxide-core/src/types/chat.rs @@ -1,6 +1,6 @@ use serde::{Deserialize, Serialize}; -use crate::types::{ChatId, ChatLocation, ChatPermissions, ChatPhoto, Message, True}; +use crate::types::{ChatId, ChatLocation, ChatPermissions, ChatPhoto, Message, True, User}; /// This object represents a chat. /// @@ -493,6 +493,23 @@ impl Chat { _ => None, } } + + /// Returns all users that are "contained" in this `Chat` + /// structure. + /// + /// This might be useful to track information about users. + /// + /// Note that this function can return duplicate users. + pub fn mentioned_users(&self) -> impl Iterator { + crate::util::flatten(self.pinned_message.as_ref().map(|m| m.mentioned_users())) + } + + /// `{Message, Chat}::mentioned_users` are mutually recursive, as such we + /// can't use `->impl Iterator` everywhere, as it would make an infinite + /// type. So we need to box somewhere. + pub(crate) fn mentioned_users_rec(&self) -> impl Iterator { + crate::util::flatten(self.pinned_message.as_ref().map(|m| m.mentioned_users_rec())) + } } mod serde_helper { diff --git a/crates/teloxide-core/src/types/chat_join_request.rs b/crates/teloxide-core/src/types/chat_join_request.rs index ec5210d7..455d922e 100644 --- a/crates/teloxide-core/src/types/chat_join_request.rs +++ b/crates/teloxide-core/src/types/chat_join_request.rs @@ -18,3 +18,15 @@ pub struct ChatJoinRequest { /// Chat invite link that was used by the user to send the join request pub invite_link: Option, } + +impl ChatJoinRequest { + /// Returns all users that are "contained" in this `ChatJoinRequest` + /// structure. + /// + /// This might be useful to track information about users. + /// + /// Note that this function can return duplicate users. + pub fn mentioned_users(&self) -> impl Iterator { + std::iter::once(&self.from).chain(self.chat.mentioned_users()) + } +} diff --git a/crates/teloxide-core/src/types/chat_member_updated.rs b/crates/teloxide-core/src/types/chat_member_updated.rs index 86094d2e..61ec402d 100644 --- a/crates/teloxide-core/src/types/chat_member_updated.rs +++ b/crates/teloxide-core/src/types/chat_member_updated.rs @@ -20,3 +20,21 @@ pub struct ChatMemberUpdated { /// joining by invite link events only. pub invite_link: Option, } + +impl ChatMemberUpdated { + /// Returns all users that are "contained" in this `ChatMemberUpdated` + /// structure. + /// + /// This might be useful to track information about users. + /// + /// Note that this function can return duplicate users. + pub fn mentioned_users(&self) -> impl Iterator { + [ + &self.from, + /* ignore `old_chat_member.user`, it should always be the same as the new one */ + &self.new_chat_member.user, + ] + .into_iter() + .chain(self.chat.mentioned_users()) + } +} diff --git a/crates/teloxide-core/src/types/game.rs b/crates/teloxide-core/src/types/game.rs index fa9c5f33..5fdc06a4 100644 --- a/crates/teloxide-core/src/types/game.rs +++ b/crates/teloxide-core/src/types/game.rs @@ -1,6 +1,6 @@ use serde::{Deserialize, Serialize}; -use crate::types::{Animation, MessageEntity, PhotoSize}; +use crate::types::{Animation, MessageEntity, PhotoSize, User}; /// This object represents a game. /// @@ -39,3 +39,17 @@ pub struct Game { /// [@Botfather]: https://t.me/botfather pub animation: Option, } + +impl Game { + /// Returns all users that are "contained" in this `Game` + /// structure. + /// + /// This might be useful to track information about users. + /// + /// Note that this function can return duplicate users. + pub fn mentioned_users(&self) -> impl Iterator { + use crate::util::{flatten, mentioned_users_from_entities}; + + flatten(self.text_entities.as_deref().map(mentioned_users_from_entities)) + } +} diff --git a/crates/teloxide-core/src/types/message.rs b/crates/teloxide-core/src/types/message.rs index 0138be1a..2c6403d8 100644 --- a/crates/teloxide-core/src/types/message.rs +++ b/crates/teloxide-core/src/types/message.rs @@ -1406,6 +1406,42 @@ impl Message { pub fn parse_caption_entities(&self) -> Option>> { self.caption().zip(self.caption_entities()).map(|(t, e)| MessageEntityRef::parse(t, e)) } + + /// Returns all users that are "contained" in this `Message` structure. + /// + /// This might be useful to track information about users. + /// + /// Note that this function may return quite a few users as it scans + /// replies, pinned messages, message entities and more. Also note that this + /// function can return duplicate users. + pub fn mentioned_users(&self) -> impl Iterator { + use crate::util::{flatten, mentioned_users_from_entities}; + + // Lets just hope we didn't forget something here... + + self.from() + .into_iter() + .chain(self.via_bot.as_ref()) + .chain(self.chat.mentioned_users_rec()) + .chain(flatten(self.reply_to_message().map(Self::mentioned_users_rec))) + .chain(flatten(self.new_chat_members())) + .chain(self.left_chat_member()) + .chain(self.forward_from_user()) + .chain(flatten(self.forward_from_chat().map(Chat::mentioned_users_rec))) + .chain(flatten(self.game().map(Game::mentioned_users))) + .chain(flatten(self.entities().map(mentioned_users_from_entities))) + .chain(flatten(self.caption_entities().map(mentioned_users_from_entities))) + .chain(flatten(self.poll().map(Poll::mentioned_users))) + .chain(flatten(self.proximity_alert_triggered().map(|a| [&a.traveler, &a.watcher]))) + .chain(flatten(self.video_chat_participants_invited().and_then(|i| i.users.as_deref()))) + } + + /// `Message::mentioned_users` is recursive (due to replies), as such we + /// can't use `->impl Iterator` everywhere, as it would make an infinite + /// type. So we need to box somewhere. + pub(crate) fn mentioned_users_rec(&self) -> Box + Send + Sync + '_> { + Box::new(self.mentioned_users()) + } } #[cfg(test)] diff --git a/crates/teloxide-core/src/types/poll.rs b/crates/teloxide-core/src/types/poll.rs index 4a35a7d5..7725c961 100644 --- a/crates/teloxide-core/src/types/poll.rs +++ b/crates/teloxide-core/src/types/poll.rs @@ -1,4 +1,4 @@ -use crate::types::{MessageEntity, PollType}; +use crate::types::{MessageEntity, PollType, User}; use chrono::{DateTime, Utc}; use serde::{Deserialize, Serialize}; @@ -67,6 +67,20 @@ pub struct PollOption { pub voter_count: i32, } +impl Poll { + /// Returns all users that are "contained" in this `Poll` + /// structure. + /// + /// This might be useful to track information about users. + /// + /// Note that this function can return duplicate users. + pub fn mentioned_users(&self) -> impl Iterator { + use crate::util::{flatten, mentioned_users_from_entities}; + + flatten(self.explanation_entities.as_deref().map(mentioned_users_from_entities)) + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/crates/teloxide-core/src/types/update.rs b/crates/teloxide-core/src/types/update.rs index e63a42f5..67c35862 100644 --- a/crates/teloxide-core/src/types/update.rs +++ b/crates/teloxide-core/src/types/update.rs @@ -59,6 +59,48 @@ impl Update { Some(from) } + /// Returns all users that are "contained" in this `Update` structure. + /// + /// This might be useful to track information about users. + /// + /// Note that this function may return quite a few users as it scans + /// replies, pinned messages, message entities, "via bot" fields and more. + /// Also note that this function can return duplicate users. + pub fn mentioned_users(&self) -> impl Iterator { + use either::Either::{Left, Right}; + use std::iter::{empty, once}; + + let i0 = Left; + let i1 = |x| Right(Left(x)); + let i2 = |x| Right(Right(Left(x))); + let i3 = |x| Right(Right(Right(Left(x)))); + let i4 = |x| Right(Right(Right(Right(Left(x))))); + let i5 = |x| Right(Right(Right(Right(Right(Left(x)))))); + let i6 = |x| Right(Right(Right(Right(Right(Right(x)))))); + + match &self.kind { + UpdateKind::Message(message) + | UpdateKind::EditedMessage(message) + | UpdateKind::ChannelPost(message) + | UpdateKind::EditedChannelPost(message) => i0(message.mentioned_users()), + + UpdateKind::InlineQuery(query) => i1(once(&query.from)), + UpdateKind::ChosenInlineResult(query) => i1(once(&query.from)), + UpdateKind::CallbackQuery(query) => i2(query.mentioned_users()), + UpdateKind::ShippingQuery(query) => i1(once(&query.from)), + UpdateKind::PreCheckoutQuery(query) => i1(once(&query.from)), + UpdateKind::Poll(poll) => i3(poll.mentioned_users()), + + UpdateKind::PollAnswer(answer) => i1(once(&answer.user)), + + UpdateKind::MyChatMember(member) | UpdateKind::ChatMember(member) => { + i4(member.mentioned_users()) + } + UpdateKind::ChatJoinRequest(request) => i5(request.mentioned_users()), + UpdateKind::Error(_) => i6(empty()), + } + } + /// Returns the chat in which is update has happened, if any. #[must_use] pub fn chat(&self) -> Option<&Chat> { diff --git a/crates/teloxide-core/src/util.rs b/crates/teloxide-core/src/util.rs new file mode 100644 index 00000000..a917428b --- /dev/null +++ b/crates/teloxide-core/src/util.rs @@ -0,0 +1,56 @@ +use crate::types::{MessageEntity, User}; + +/// Converts an optional iterator to a flattened iterator. +pub(crate) fn flatten(opt: Option) -> impl Iterator +where + I: IntoIterator, +{ + struct Flat(Option); + + impl Iterator for Flat + where + I: Iterator, + { + type Item = I::Item; + + fn next(&mut self) -> Option { + self.0.as_mut()?.next() + } + + fn size_hint(&self) -> (usize, Option) { + match &self.0 { + None => (0, Some(0)), + Some(i) => i.size_hint(), + } + } + } + + Flat(opt.map(<_>::into_iter)) +} + +pub(crate) fn mentioned_users_from_entities( + entities: &[MessageEntity], +) -> impl Iterator { + use crate::types::MessageEntityKind::*; + + entities.iter().filter_map(|entity| match &entity.kind { + TextMention { user } => Some(user), + + Mention + | Hashtag + | Cashtag + | BotCommand + | Url + | Email + | PhoneNumber + | Bold + | Italic + | Underline + | Strikethrough + | Spoiler + | Code + | Pre { language: _ } + | TextLink { url: _ } + | CustomEmoji { custom_emoji_id: _ } => None, + }) +}