From ad88db96b6a17c25a83d99228017c84b985433ac Mon Sep 17 00:00:00 2001
From: Hirrolot <hirrolot@gmail.com>
Date: Sun, 23 Jan 2022 20:38:12 +0600
Subject: [PATCH] Add `UpdateFilterExt`

---
 src/dispatching2/dispatcher.rs | 16 -------------
 src/dispatching2/filter_ext.rs | 44 +++++++++++++++++++++++++++-------
 src/dispatching2/mod.rs        |  2 +-
 3 files changed, 37 insertions(+), 25 deletions(-)

diff --git a/src/dispatching2/dispatcher.rs b/src/dispatching2/dispatcher.rs
index fa12328e..b3dac9dc 100644
--- a/src/dispatching2/dispatcher.rs
+++ b/src/dispatching2/dispatcher.rs
@@ -264,20 +264,4 @@ where
     pub fn dependencies(self, dependencies: DependencyMap) -> Self {
         Dispatcher { dependencies, ..self }
     }
-
-    define_handlers! {
-        (messages_handler, Message),
-        (edited_messages_handler, EditedMessage),
-        (channel_posts_handler, ChannelPost),
-        (edited_channel_posts_handler, EditedChannelPost),
-        (inline_queries_handler, InlineQuery),
-        (chosen_inline_results_handler, ChosenInlineResult),
-        (callback_queries_handler, CallbackQuery),
-        (shipping_queries_handler, ShippingQuery),
-        (pre_checkout_queries_handler, PreCheckoutQuery),
-        (polls_handler, Poll),
-        (poll_answers_handler, PollAnswer),
-        (my_chat_members_handler, MyChatMember),
-        (chat_members_handler, ChatMember),
-    }
 }
diff --git a/src/dispatching2/filter_ext.rs b/src/dispatching2/filter_ext.rs
index 4374912a..dcd6a7bd 100644
--- a/src/dispatching2/filter_ext.rs
+++ b/src/dispatching2/filter_ext.rs
@@ -1,8 +1,8 @@
 use dptree::{di::DependencyMap, Handler};
-use teloxide_core::types::Message;
+use teloxide_core::types::{Message, Update, UpdateKind};
 
 macro_rules! define_ext {
-    ($ext_name:ident, $for_ty:ty => $( ($func:ident, $arg_ty:ty, $get_func:expr) ,)*) => {
+    ($ext_name:ident, $for_ty:ty => $( ($func:ident, $arg_ty:ty, $proj_fn:expr) ,)*) => {
         pub trait $ext_name<Out> {
             $( define_ext!(@sig $func, $arg_ty); )*
         }
@@ -11,16 +11,18 @@ macro_rules! define_ext {
         where
             Out: Send + Sync + 'static,
         {
-            $( define_ext!(@impl $for_ty, $func, $arg_ty, $get_func); )*
+            $( define_ext!(@impl $for_ty, $func, $arg_ty, $proj_fn); )*
         }
     };
+
     (@sig $func:ident, $arg_ty:ty) => {
         fn $func<F, Fut>() -> Handler<'static, DependencyMap, Out>;
     };
-    (@impl $for_ty:ty, $func:ident, $arg_ty:ty, $get_func:expr) => {
+
+    (@impl $for_ty:ty, $func:ident, $arg_ty:ty, $proj_fn:expr) => {
         fn $func<F, Fut>() -> Handler<'static, DependencyMap, Out> {
-            dptree::filter_map(move |msg: $for_ty| {
-                let result = $get_func(&msg).map(ToOwned::to_owned);
+            dptree::filter_map(move |input: $for_ty| {
+                let result = $proj_fn(&input).map(ToOwned::to_owned);
                 async move { result }
             })
         }
@@ -28,7 +30,7 @@ macro_rules! define_ext {
 }
 
 // May be expanded in the future.
-define_ext!(
+define_ext! {
     MessageFilterExt, Message =>
     (filter_from, types::User, Message::from),
     (filter_animation, types::Animation, Message::animation),
@@ -46,4 +48,30 @@ define_ext!(
     (filter_left_chat_member, types::User, Message::left_chat_member),
     (filter_pinned, Message, Message::pinned_message),
     (filter_dice, types::Dice, Message::dice),
-);
+}
+
+macro_rules! kind {
+    ($kind:ident) => {
+        |update: &Update| match update.kind {
+            UpdateKind::$kind(x) => Some(&x),
+            _ => None,
+        }
+    };
+}
+
+define_ext! {
+    UpdateFilterExt, Update =>
+    (filter_message, types::Message, kind![Message]),
+    (filter_edited_message, types::EditedMessage, kind![EditedMessage]),
+    (filter_channel_post, types::ChannelPost, kind![ChannelPost]),
+    (filter_edited_channel_post, types::EditedChannelPost, kind![EditedChannelPost]),
+    (filter_inline_query, types::InlineQuery, kind![InlineQuery]),
+    (filter_chosen_inline_result, types::ChosenInlineResult, kind![ChosenInlineResult]),
+    (filter_callback_query, types::CallbackQuery, kind![CallbackQuery]),
+    (filter_shipping_query, types::ShippingQuery, kind![ShippingQuery]),
+    (filter_pre_checkout_query, types::PreCheckoutQuery, kind![PreCheckoutQuery]),
+    (filter_poll, types::Poll, kind![Poll]),
+    (filter_poll_answer, types::PollAnswer, kind![PollAnswer]),
+    (filter_my_chat_member, types::MyChatMember, kind![MyChatMember]),
+    (filter_chat_member, types::ChatMember, kind![ChatMember]),
+}
diff --git a/src/dispatching2/mod.rs b/src/dispatching2/mod.rs
index 5a8a2fd6..13cdf437 100644
--- a/src/dispatching2/mod.rs
+++ b/src/dispatching2/mod.rs
@@ -7,6 +7,6 @@ mod handler_ext;
 mod handler_factory;
 
 pub use dispatcher::{DefaultHandler, Dispatcher, UpdateHandler};
-pub use filter_ext::MessageFilterExt;
+pub use filter_ext::{MessageFilterExt, UpdateFilterExt};
 pub use handler_ext::HandlerExt;
 pub use handler_factory::HandlerFactory;