From 481f0db129bdff4538d876165334da3d83da32a1 Mon Sep 17 00:00:00 2001
From: teidesu <me@tei.su>
Date: Sun, 20 Jun 2021 02:52:19 +0300
Subject: [PATCH] fix(dispatcher): provide state type information to filters

---
 packages/dispatcher/scripts/generate.js | 19 ++++++-
 packages/dispatcher/src/dispatcher.ts   | 74 ++++++++++++++++++++++---
 2 files changed, 82 insertions(+), 11 deletions(-)

diff --git a/packages/dispatcher/scripts/generate.js b/packages/dispatcher/scripts/generate.js
index bb1b7536..facacf79 100644
--- a/packages/dispatcher/scripts/generate.js
+++ b/packages/dispatcher/scripts/generate.js
@@ -230,7 +230,22 @@ function generateDispatcher() {
      * @param group  Handler group index
      * @internal
      */
-    on${type.handlerTypeName}(handler: ${type.handlerTypeName}Handler${type.state ? `<${type.updateType}, State extends never ? never : MessageState<State, SceneName>>` : ''}['callback'], group?: number): void
+    on${type.handlerTypeName}(handler: ${type.handlerTypeName}Handler${type.state ? `<${type.updateType}, State extends never ? never : UpdateState<State, SceneName>>` : ''}['callback'], group?: number): void
+
+${type.state ? `
+    /**
+     * Register ${toSentence(type)} with a filter
+     *
+     * @param filter  Update filter
+     * @param handler  ${toSentence(type, 'full')}
+     * @param group  Handler group index
+     */
+    on${type.handlerTypeName}<Mod>(
+        filter: UpdateFilter<${type.updateType}, Mod, State>,
+        handler: ${type.handlerTypeName}Handler<filters.Modify<${type.updateType}, Mod>, State extends never ? never : UpdateState<State, SceneName>>['callback'],
+        group?: number
+    ): void
+    ` : ''}
 
     /**
      * Register ${toSentence(type)} with a filter
@@ -241,7 +256,7 @@ function generateDispatcher() {
      */
     on${type.handlerTypeName}<Mod>(
         filter: UpdateFilter<${type.updateType}, Mod>,
-        handler: ${type.handlerTypeName}Handler<filters.Modify<${type.updateType}, Mod>${type.state ? ', State extends never ? never : MessageState<State, SceneName>' : ''}>['callback'],
+        handler: ${type.handlerTypeName}Handler<filters.Modify<${type.updateType}, Mod>${type.state ? ', State extends never ? never : UpdateState<State, SceneName>' : ''}>['callback'],
         group?: number
     ): void
 
diff --git a/packages/dispatcher/src/dispatcher.ts b/packages/dispatcher/src/dispatcher.ts
index 6c28c0bd..d6dc615b 100644
--- a/packages/dispatcher/src/dispatcher.ts
+++ b/packages/dispatcher/src/dispatcher.ts
@@ -25,9 +25,9 @@ import {
     PollVoteHandler,
     UserStatusUpdateHandler,
     UserTypingHandler,
-    UpdateInfoForError,
 } from './handler'
 // end-codegen-imports
+import { UpdateInfoForError } from './handler'
 import { filters, UpdateFilter } from './filters'
 import { handlers } from './builders'
 import { ChatMemberUpdate } from './updates'
@@ -461,14 +461,20 @@ export class Dispatcher<State = never, SceneName extends string = string> {
                                 return
                             case 'scene': {
                                 if (!parsedState)
-                                    throw new MtCuteArgumentError('Cannot use ToScene without state')
+                                    throw new MtCuteArgumentError(
+                                        'Cannot use ToScene without state'
+                                    )
 
                                 const scene = parsedState['_scene']
 
                                 if (!scene)
-                                    throw new MtCuteArgumentError('Cannot use ToScene without entering a scene')
+                                    throw new MtCuteArgumentError(
+                                        'Cannot use ToScene without entering a scene'
+                                    )
 
-                                return this._scenes[scene]._dispatchUpdateNowImpl(
+                                return this._scenes[
+                                    scene
+                                ]._dispatchUpdateNowImpl(
                                     update,
                                     users,
                                     chats,
@@ -978,9 +984,13 @@ export class Dispatcher<State = never, SceneName extends string = string> {
      * This will load the state for the given object
      * ignoring local custom storage, key delegate and scene scope.
      */
-    getGlobalState<T>(object: Parameters<StateKeyDelegate>[0]): Promise<UpdateState<T, SceneName>> {
+    getGlobalState<T>(
+        object: Parameters<StateKeyDelegate>[0]
+    ): Promise<UpdateState<T, SceneName>> {
         if (!this._parent) {
-            throw new MtCuteArgumentError('This dispatcher does not have a parent')
+            throw new MtCuteArgumentError(
+                'This dispatcher does not have a parent'
+            )
         }
 
         return Promise.resolve(this._stateKeyDelegate!(object)).then((key) => {
@@ -1071,9 +1081,23 @@ export class Dispatcher<State = never, SceneName extends string = string> {
         filter: UpdateFilter<Message, Mod, State>,
         handler: NewMessageHandler<
             filters.Modify<Message, Mod>,
-            State extends never
-                ? never
-                : UpdateState<State, SceneName>
+            State extends never ? never : UpdateState<State, SceneName>
+        >['callback'],
+        group?: number
+    ): void
+
+    /**
+     * Register a new message handler with a filter
+     *
+     * @param filter  Update filter
+     * @param handler  New message handler
+     * @param group  Handler group index
+     */
+    onNewMessage<Mod>(
+        filter: UpdateFilter<Message, Mod>,
+        handler: NewMessageHandler<
+            filters.Modify<Message, Mod>,
+            State extends never ? never : UpdateState<State, SceneName>
         >['callback'],
         group?: number
     ): void
@@ -1098,6 +1122,22 @@ export class Dispatcher<State = never, SceneName extends string = string> {
         group?: number
     ): void
 
+    /**
+     * Register an edit message handler with a filter
+     *
+     * @param filter  Update filter
+     * @param handler  Edit message handler
+     * @param group  Handler group index
+     */
+    onEditMessage<Mod>(
+        filter: UpdateFilter<Message, Mod, State>,
+        handler: EditMessageHandler<
+            filters.Modify<Message, Mod>,
+            State extends never ? never : UpdateState<State, SceneName>
+        >['callback'],
+        group?: number
+    ): void
+
     /**
      * Register an edit message handler with a filter
      *
@@ -1259,6 +1299,22 @@ export class Dispatcher<State = never, SceneName extends string = string> {
         group?: number
     ): void
 
+    /**
+     * Register a callback query handler with a filter
+     *
+     * @param filter  Update filter
+     * @param handler  Callback query handler
+     * @param group  Handler group index
+     */
+    onCallbackQuery<Mod>(
+        filter: UpdateFilter<CallbackQuery, Mod, State>,
+        handler: CallbackQueryHandler<
+            filters.Modify<CallbackQuery, Mod>,
+            State extends never ? never : UpdateState<State, SceneName>
+        >['callback'],
+        group?: number
+    ): void
+
     /**
      * Register a callback query handler with a filter
      *