ronami/dispatching/
handler_ext.rs

1use crate::{
2    dispatching::{
3        dialogue::{GetChatId, Storage},
4        DpHandlerDescription,
5    },
6    types::{Me, Message},
7    utils::command::BotCommands,
8};
9use dptree::Handler;
10
11use std::fmt::Debug;
12
13/// Extension methods for working with `dptree` handlers.
14pub trait HandlerExt<Output> {
15    /// Returns a handler that accepts a parsed command `C`.
16    ///
17    /// ## Dependency requirements
18    ///
19    ///  - [`crate::types::Message`]
20    ///  - [`crate::types::Me`]
21    #[must_use]
22    fn filter_command<C>(self) -> Self
23    where
24        C: BotCommands + Send + Sync + 'static;
25
26    /// Returns a handler that accepts a parsed command `C` if the command
27    /// contains a bot mention, for example `/start@my_bot`.
28    ///
29    /// ## Dependency requirements
30    ///
31    ///  - [`crate::types::Message`]
32    ///  - [`crate::types::Me`]
33    #[must_use]
34    fn filter_mention_command<C>(self) -> Self
35    where
36        C: BotCommands + Send + Sync + 'static;
37
38    /// Passes [`Dialogue<D, S>`] and `D` as handler dependencies.
39    ///
40    /// It does so by the following steps:
41    ///
42    ///  1. If an incoming update has no chat ID ([`GetChatId::chat_id`] returns
43    ///     `None`), the rest of the chain will not be executed. Otherwise,
44    ///     passes `Dialogue::new(storage, chat_id)` forwards.
45    ///  2. If [`Dialogue::get_or_default`] on the passed dialogue returns `Ok`,
46    ///     passes the dialogue state forwards. Otherwise, logs an error and the
47    ///     rest of the chain is not executed.
48    ///
49    /// If `RONAMI_DIALOGUE_BEHAVIOUR` environmental variable exists and is
50    /// equal to "default", this function will not panic if it can't get the
51    /// dialogue (if, for example, the state enum was updated). Setting the
52    /// value to "panic" will return the initial behaviour.
53    ///
54    /// ## Dependency requirements
55    ///
56    ///  - `Arc<S>`
57    ///  - `Upd`
58    ///
59    /// [`Dialogue<D, S>`]: super::dialogue::Dialogue
60    /// [`Dialogue::get_or_default`]: super::dialogue::Dialogue::get_or_default
61    #[must_use]
62    fn enter_dialogue<Upd, S, D>(self) -> Self
63    where
64        S: Storage<D> + ?Sized + Send + Sync + 'static,
65        <S as Storage<D>>::Error: Debug + Send,
66        D: Default + Clone + Send + Sync + 'static,
67        Upd: GetChatId + Clone + Send + Sync + 'static;
68}
69
70impl<Output> HandlerExt<Output> for Handler<'static, Output, DpHandlerDescription>
71where
72    Output: Send + Sync + 'static,
73{
74    fn filter_command<C>(self) -> Self
75    where
76        C: BotCommands + Send + Sync + 'static,
77    {
78        self.chain(filter_command::<C, Output>())
79    }
80
81    fn filter_mention_command<C>(self) -> Self
82    where
83        C: BotCommands + Send + Sync + 'static,
84    {
85        self.chain(filter_mention_command::<C, Output>())
86    }
87
88    fn enter_dialogue<Upd, S, D>(self) -> Self
89    where
90        S: Storage<D> + ?Sized + Send + Sync + 'static,
91        <S as Storage<D>>::Error: Debug + Send,
92        D: Default + Clone + Send + Sync + 'static,
93        Upd: GetChatId + Clone + Send + Sync + 'static,
94    {
95        self.chain(super::dialogue::enter::<Upd, S, D, Output>())
96    }
97}
98
99/// Returns a handler that accepts a parsed command `C`.
100///
101/// A call to this function is the same as `dptree::entry().filter_command()`.
102///
103/// See [`HandlerExt::filter_command`].
104///
105/// ## Dependency requirements
106///
107///  - [`crate::types::Message`]
108///  - [`crate::types::Me`]
109#[must_use]
110pub fn filter_command<C, Output>() -> Handler<'static, Output, DpHandlerDescription>
111where
112    C: BotCommands + Send + Sync + 'static,
113    Output: Send + Sync + 'static,
114{
115    dptree::filter_map(move |message: Message, me: Me| {
116        let bot_name = me.user.username.expect("Bots must have a username");
117        message.text().or_else(|| message.caption()).and_then(|text| C::parse(text, &bot_name).ok())
118    })
119}
120
121/// Returns a handler that accepts a parsed command `C` if the command
122/// contains a bot mention, for example `/start@my_bot`.
123///
124/// A call to this function is the same as
125/// `dptree::entry().filter_mention_command()`.
126///
127/// See [`HandlerExt::filter_mention_command`].
128///
129/// ## Dependency requirements
130///
131///  - [`crate::types::Message`]
132///  - [`crate::types::Me`]
133#[must_use]
134pub fn filter_mention_command<C, Output>() -> Handler<'static, Output, DpHandlerDescription>
135where
136    C: BotCommands + Send + Sync + 'static,
137    Output: Send + Sync + 'static,
138{
139    dptree::filter_map(move |message: Message, me: Me| {
140        let bot_name = me.user.username.expect("Bots must have a username");
141
142        let text_or_caption = message.text().or_else(|| message.caption());
143        let command = text_or_caption.and_then(|text| C::parse(text, &bot_name).ok());
144        // If the parsing succeeds with a bot_name,
145        // but fails without - there is a mention
146        let is_username_required =
147            text_or_caption.and_then(|text| C::parse(text, "").ok()).is_none();
148
149        if !is_username_required {
150            return None;
151        }
152        command
153    })
154}
155
156#[cfg(test)]
157#[cfg(feature = "macros")]
158mod tests {
159    use crate::{self as ronami, dispatching::UpdateFilterExt, utils::command::BotCommands};
160    use chrono::DateTime;
161    use dptree::deps;
162    use ronami_core::types::{
163        Chat, ChatId, ChatKind, ChatPrivate, LinkPreviewOptions, Me, MediaKind, MediaText, Message,
164        MessageCommon, MessageId, MessageKind, Update, UpdateId, UpdateKind, User, UserId,
165    };
166
167    use super::HandlerExt;
168
169    #[derive(BotCommands, Clone)]
170    #[command(rename_rule = "lowercase")]
171    enum Cmd {
172        Test,
173    }
174
175    fn make_update(text: String) -> Update {
176        let timestamp = 1_569_518_829;
177        let date = DateTime::from_timestamp(timestamp, 0).unwrap();
178        Update {
179            id: UpdateId(326_170_274),
180            kind: UpdateKind::Message(Message {
181                via_bot: None,
182                id: MessageId(5042),
183                thread_id: None,
184                from: Some(User {
185                    id: UserId(109_998_024),
186                    is_bot: false,
187                    first_name: String::from("Laster"),
188                    last_name: None,
189                    username: Some(String::from("laster_alex")),
190                    language_code: Some(String::from("en")),
191                    is_premium: false,
192                    added_to_attachment_menu: false,
193                    can_manage_bots: false,
194                    supports_guest_queries: false,
195                }),
196                sender_chat: None,
197                is_topic_message: false,
198                is_paid_post: false,
199                suggested_post_info: None,
200                sender_business_bot: None,
201                date,
202                chat: Chat {
203                    id: ChatId(109_998_024),
204                    kind: ChatKind::Private(ChatPrivate {
205                        username: Some(String::from("Laster")),
206                        first_name: Some(String::from("laster_alex")),
207                        last_name: None,
208                    }),
209                },
210                direct_messages_topic: None,
211                kind: MessageKind::Common(MessageCommon {
212                    reply_to_message: None,
213                    forward_origin: None,
214                    external_reply: None,
215                    quote: None,
216                    edit_date: None,
217                    media_kind: MediaKind::Text(MediaText {
218                        text,
219                        entities: vec![],
220                        link_preview_options: Some(LinkPreviewOptions {
221                            is_disabled: true,
222                            url: None,
223                            prefer_small_media: false,
224                            prefer_large_media: false,
225                            show_above_text: false,
226                        }),
227                    }),
228                    reply_markup: None,
229                    author_signature: None,
230                    paid_star_count: None,
231                    effect_id: None,
232                    is_automatic_forward: false,
233                    has_protected_content: false,
234                    reply_to_checklist_task_id: None,
235                    reply_to_story: None,
236                    reply_to_poll_option_id: None,
237                    sender_tag: None,
238                    guest_query_id: None,
239                    guest_bot_caller_user: None,
240                    guest_bot_caller_chat: None,
241                    live_photo: None,
242                    sender_boost_count: None,
243                    is_from_offline: false,
244                    business_connection_id: None,
245                }),
246            }),
247        }
248    }
249
250    fn make_me() -> Me {
251        Me {
252            user: User {
253                id: UserId(42),
254                is_bot: true,
255                first_name: "First".to_owned(),
256                last_name: None,
257                username: Some("SomethingSomethingBot".to_owned()),
258                language_code: None,
259                is_premium: false,
260                added_to_attachment_menu: false,
261                can_manage_bots: false,
262                supports_guest_queries: false,
263            },
264            can_join_groups: false,
265            allows_users_to_create_topics: false,
266            can_read_all_group_messages: false,
267            supports_inline_queries: false,
268            can_connect_to_business: false,
269            has_main_web_app: false,
270            has_topics_enabled: false,
271        }
272    }
273
274    #[tokio::test]
275    async fn test_filter_command() {
276        let h = dptree::entry()
277            .branch(Update::filter_message().filter_command::<Cmd>().endpoint(|| async {}));
278        let me = make_me();
279
280        let update = make_update("/test@".to_owned() + me.username());
281        let result = h.dispatch(deps![update, me.clone()]).await;
282        assert!(result.is_break());
283
284        let update = make_update("/test@".to_owned() + "SomeOtherBot");
285        let result = h.dispatch(deps![update, me.clone()]).await;
286        assert!(result.is_continue());
287
288        let update = make_update("/test".to_owned());
289        let result = h.dispatch(deps![update, me.clone()]).await;
290        assert!(result.is_break());
291    }
292
293    #[tokio::test]
294    async fn test_filter_mention_command() {
295        let h = dptree::entry()
296            .branch(Update::filter_message().filter_mention_command::<Cmd>().endpoint(|| async {}));
297        let me = make_me();
298
299        let update = make_update("/test@".to_owned() + me.username());
300        let result = h.dispatch(deps![update, me.clone()]).await;
301        assert!(result.is_break());
302
303        let update = make_update("/test@".to_owned() + "SomeOtherBot");
304        let result = h.dispatch(deps![update, me.clone()]).await;
305        assert!(result.is_continue());
306
307        let update = make_update("/test".to_owned());
308        let result = h.dispatch(deps![update, me.clone()]).await;
309        assert!(result.is_continue());
310    }
311}
ronami/dispatching/handler_ext.rs

ronami/dispatching/
handler_ext.rs
