teloxide_ng/dispatching/
handler_ext.rs

1use crate::{
2    dispatching::{
3        DpHandlerDescription,
4        dialogue::{GetChatId, Storage},
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 `TELOXIDE_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 teloxide_ng, dispatching::UpdateFilterExt, utils::command::BotCommands};
160    use chrono::DateTime;
161    use dptree::deps;
162    use teloxide_core_ng::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                    has_topics_enabled: false,
194                }),
195                sender_chat: None,
196                is_topic_message: false,
197                sender_business_bot: None,
198                date,
199                chat: Chat {
200                    id: ChatId(109_998_024),
201                    kind: ChatKind::Private(ChatPrivate {
202                        username: Some(String::from("Laster")),
203                        first_name: Some(String::from("laster_alex")),
204                        last_name: None,
205                    }),
206                },
207                kind: MessageKind::Common(MessageCommon {
208                    reply_to_message: None,
209                    forward_origin: None,
210                    external_reply: None,
211                    quote: None,
212                    edit_date: None,
213                    media_kind: MediaKind::Text(MediaText {
214                        text,
215                        entities: vec![],
216                        link_preview_options: Some(LinkPreviewOptions {
217                            is_disabled: true,
218                            url: None,
219                            prefer_small_media: false,
220                            prefer_large_media: false,
221                            show_above_text: false,
222                        }),
223                    }),
224                    reply_markup: None,
225                    author_signature: None,
226                    paid_star_count: None,
227                    effect_id: None,
228                    is_automatic_forward: false,
229                    has_protected_content: false,
230                    reply_to_story: None,
231                    sender_boost_count: None,
232                    is_from_offline: false,
233                    business_connection_id: None,
234                }),
235            }),
236        }
237    }
238
239    fn make_me() -> Me {
240        Me {
241            user: User {
242                id: UserId(42),
243                is_bot: true,
244                first_name: "First".to_owned(),
245                last_name: None,
246                username: Some("SomethingSomethingBot".to_owned()),
247                language_code: None,
248                is_premium: false,
249                added_to_attachment_menu: false,
250                has_topics_enabled: false,
251            },
252            can_join_groups: false,
253            can_read_all_group_messages: false,
254            supports_inline_queries: false,
255            can_connect_to_business: false,
256            has_main_web_app: false,
257        }
258    }
259
260    #[tokio::test]
261    async fn test_filter_command() {
262        let h = dptree::entry()
263            .branch(Update::filter_message().filter_command::<Cmd>().endpoint(|| async {}));
264        let me = make_me();
265
266        let update = make_update("/test@".to_owned() + me.username());
267        let result = h.dispatch(deps![update, me.clone()]).await;
268        assert!(result.is_break());
269
270        let update = make_update("/test@".to_owned() + "SomeOtherBot");
271        let result = h.dispatch(deps![update, me.clone()]).await;
272        assert!(result.is_continue());
273
274        let update = make_update("/test".to_owned());
275        let result = h.dispatch(deps![update, me.clone()]).await;
276        assert!(result.is_break());
277    }
278
279    #[tokio::test]
280    async fn test_filter_mention_command() {
281        let h = dptree::entry()
282            .branch(Update::filter_message().filter_mention_command::<Cmd>().endpoint(|| async {}));
283        let me = make_me();
284
285        let update = make_update("/test@".to_owned() + me.username());
286        let result = h.dispatch(deps![update, me.clone()]).await;
287        assert!(result.is_break());
288
289        let update = make_update("/test@".to_owned() + "SomeOtherBot");
290        let result = h.dispatch(deps![update, me.clone()]).await;
291        assert!(result.is_continue());
292
293        let update = make_update("/test".to_owned());
294        let result = h.dispatch(deps![update, me.clone()]).await;
295        assert!(result.is_continue());
296    }
297}
teloxide_ng/dispatching/handler_ext.rs

teloxide_ng/dispatching/
handler_ext.rs
