From 13e4740d0a8ba911eab43628d492b40155297586 Mon Sep 17 00:00:00 2001
From: InsanusMokrassar <ovsyannikov.alexey95@gmail.com>
Date: Thu, 29 Jun 2023 13:45:08 +0600
Subject: [PATCH] improve Update.sourceChat

---
 CHANGELOG.md                                  |  3 ++
 .../utils/extensions/UpdateChatRetriever.kt   | 51 +++++++++++++++++--
 2 files changed, 50 insertions(+), 4 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 7f07b74d29..e66c33cfcf 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,9 @@
 
 ## 8.1.1
 
+* `Utils`:
+    * Improve extension `Update.sourceChat` to add opportunity to select some chats by logic different with the default
+
 ## 8.1.0
 
 **PARTIALLY BREAKING CHANGES: Exclude `.*Impl` classcasts from `ClassCastsNew`**
diff --git a/tgbotapi.utils/src/commonMain/kotlin/dev/inmo/tgbotapi/extensions/utils/extensions/UpdateChatRetriever.kt b/tgbotapi.utils/src/commonMain/kotlin/dev/inmo/tgbotapi/extensions/utils/extensions/UpdateChatRetriever.kt
index 56ec9dc233..4c7911fe71 100644
--- a/tgbotapi.utils/src/commonMain/kotlin/dev/inmo/tgbotapi/extensions/utils/extensions/UpdateChatRetriever.kt
+++ b/tgbotapi.utils/src/commonMain/kotlin/dev/inmo/tgbotapi/extensions/utils/extensions/UpdateChatRetriever.kt
@@ -3,17 +3,60 @@ package dev.inmo.tgbotapi.extensions.utils.extensions
 import dev.inmo.tgbotapi.abstracts.FromUser
 import dev.inmo.tgbotapi.abstracts.WithUser
 import dev.inmo.tgbotapi.extensions.utils.asUser
+import dev.inmo.tgbotapi.types.InlineQueries.ChosenInlineResult.BaseChosenInlineResult
+import dev.inmo.tgbotapi.types.InlineQueries.ChosenInlineResult.ChosenInlineResult
+import dev.inmo.tgbotapi.types.InlineQueries.ChosenInlineResult.LocationChosenInlineResult
+import dev.inmo.tgbotapi.types.InlineQueries.query.BaseInlineQuery
+import dev.inmo.tgbotapi.types.InlineQueries.query.LocationInlineQuery
 import dev.inmo.tgbotapi.types.chat.Chat
 import dev.inmo.tgbotapi.types.chat.User
-import dev.inmo.tgbotapi.types.update.ChatJoinRequestUpdate
+import dev.inmo.tgbotapi.types.queries.callback.*
+import dev.inmo.tgbotapi.types.update.*
 import dev.inmo.tgbotapi.types.update.abstracts.BaseMessageUpdate
 import dev.inmo.tgbotapi.types.update.abstracts.Update
 import dev.inmo.tgbotapi.utils.PreviewFeature
 
+fun CallbackQuery.sourceChat() = when (this) {
+    is InlineMessageIdDataCallbackQuery -> null
+    is MessageDataCallbackQuery -> message.chat
+    is InlineMessageIdGameShortNameCallbackQuery -> null
+    is MessageGameShortNameCallbackQuery -> message.chat
+    is UnknownCallbackQueryType -> null
+}
+
 @PreviewFeature
-fun Update.sourceChat(): Chat? = when (this) {
-    is BaseMessageUpdate -> data.chat
-    is ChatJoinRequestUpdate -> data.chat
+fun Update.sourceChat(
+    baseMessageUpdateConverter: (BaseMessageUpdate) -> Chat? = { it.data.chat },
+    chatJoinRequestUpdateConverter: (ChatJoinRequestUpdate) -> Chat? = { it.data.chat },
+    shippingQueryUpdateConverter: (ShippingQueryUpdate) -> Chat? = { it.data.from },
+    pollAnswerUpdateConverter: (PollAnswerUpdate) -> Chat? = { it.data.from },
+    preCheckoutQueryUpdateConverter: (PreCheckoutQueryUpdate) -> Chat? = { it.data.from },
+    callbackQueryUpdateConverter: (CallbackQueryUpdate) -> Chat? = { it.data.sourceChat() },
+    chosenInlineResultUpdateConverter: (ChosenInlineResultUpdate) -> Chat? = { null },
+    inlineQueryUpdateConverter: (InlineQueryUpdate) -> Chat? = { null },
+    pollUpdateConverter: (PollUpdate) -> Chat? = { null },
+    channelPostUpdateConverter: (ChannelPostUpdate) -> Chat? = { it.data.chat },
+    messageUpdateConverter: (MessageUpdate) -> Chat? = { it.data.chat },
+    editChannelPostUpdateConverter: (EditChannelPostUpdate) -> Chat? = { it.data.chat },
+    editMessageUpdateConverter: (EditMessageUpdate) -> Chat? = { it.data.chat },
+    myChatMemberUpdatedUpdateConverter: (MyChatMemberUpdatedUpdate) -> Chat? = { it.data.chat },
+    commonChatMemberUpdatedUpdateConverter: (CommonChatMemberUpdatedUpdate) -> Chat? = { it.data.chat }
+): Chat? = when (this) {
+    is BaseMessageUpdate -> baseMessageUpdateConverter(this)
+    is ChatJoinRequestUpdate -> chatJoinRequestUpdateConverter(this)
+    is ShippingQueryUpdate -> shippingQueryUpdateConverter(this)
+    is PollAnswerUpdate -> pollAnswerUpdateConverter(this)
+    is PreCheckoutQueryUpdate -> preCheckoutQueryUpdateConverter(this)
+    is CallbackQueryUpdate -> callbackQueryUpdateConverter(this)
+    is ChosenInlineResultUpdate -> chosenInlineResultUpdateConverter(this)
+    is InlineQueryUpdate -> inlineQueryUpdateConverter(this)
+    is PollUpdate -> pollUpdateConverter(this)
+    is ChannelPostUpdate -> channelPostUpdateConverter(this)
+    is MessageUpdate -> messageUpdateConverter(this)
+    is EditChannelPostUpdate -> editChannelPostUpdateConverter(this)
+    is EditMessageUpdate -> editMessageUpdateConverter(this)
+    is MyChatMemberUpdatedUpdate -> myChatMemberUpdatedUpdateConverter(this)
+    is CommonChatMemberUpdatedUpdate -> commonChatMemberUpdatedUpdateConverter(this)
     else -> {
         when (val data = data) {
             is FromUser -> data.from