From 07b9d4f2407e84541e2459ec3b30249576ca4526 Mon Sep 17 00:00:00 2001
From: lifegpc <g1710431395@gmail.com>
Date: Wed, 15 Dec 2021 13:26:21 +0800
Subject: [PATCH] add an option to delete command message

---
 tdlib.py       |  7 +++++--
 tg_user_bot.py | 14 +++++++++-----
 2 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/tdlib.py b/tdlib.py
index b597cbd..0ea2390 100644
--- a/tdlib.py
+++ b/tdlib.py
@@ -9,7 +9,7 @@ from random import random
 import sys
 from threading import Thread
 from traceback import print_exc
-from typing import List
+from typing import List, Union
 
 
 tdjson_path = find_library('tdjson')
@@ -532,8 +532,11 @@ class TdLib:
                 print(f"{re['code']} {re['message']}")
             return False
 
-    async def deleteMessages(self, chat_id: int, message_ids: List[int],
+    async def deleteMessages(self, chat_id: int,
+                             message_ids: Union[List[int], int],
                              revoke: bool = True):
+        if isinstance(message_ids, int):
+            message_ids = [message_ids]
         re = await self._send({"@type": "deleteMessages", "chat_id": chat_id,
                                "message_ids": message_ids, "revoke": revoke})
         if re['@type'] == 'ok':
diff --git a/tg_user_bot.py b/tg_user_bot.py
index 8a2236d..b79d95c 100644
--- a/tg_user_bot.py
+++ b/tg_user_bot.py
@@ -55,6 +55,7 @@ asts.add_argument('name', help="Sticker set name. Must ended with _by_<botname>"
 asts.add_argument('--id', help='Sticker set ID. Needed if name not specified.', type=int, metavar='ID', dest='id')  # noqa: E501
 asts.add_argument('-e', '--emoji', help='Emojis corresponding to the sticker', metavar='EMOJI', dest='emoji')  # noqa: E501
 asts.add_argument('-a', '--add-suffix', help='Add _by_<botname> automatically', action='store_true', dest='add_suffix')  # noqa: E501
+asts.add_argument('-d', '--delete', help='Delete command message if executed successfully.', action='store_true', dest='delete')  # noqa: E501
 
 
 def generateFileInfo(f: dict) -> str:
@@ -270,10 +271,13 @@ async def handle_add_sticker_to_set(lib: TdLib, robot: TdLib, mes: dict,
             uid = await lib.getUid()
             if uid is None:
                 raise ValueError('Can not get user ID.')
-            r = await robot.addStickerToSet(uid, r.name, st)
-            if r is None:
+            re = await robot.addStickerToSet(uid, r.name, st)
+            if re is None:
                 raise ValueError('Failed to add sticker to set')
-            re = await lib.editMessageText(mes['chat_id'], mes['id'], "Added sticker to set successfully.")  # noqa: E501
+            if r.delete:
+                re = await lib.deleteMessages(mes['chat_id'], mes['id'])
+            else:
+                re = await lib.editMessageText(mes['chat_id'], mes['id'], "Added sticker to set successfully.")  # noqa: E501
         except ValueError as e:
             if len(e.args) == 0:
                 re = await lib.editMessageText(mes['chat_id'], mes['id'], "Unknown error.")  # noqa: E501
@@ -281,8 +285,8 @@ async def handle_add_sticker_to_set(lib: TdLib, robot: TdLib, mes: dict,
                 re = await lib.editMessageText(mes['chat_id'], mes['id'], f"```\n{asts.format_usage()}{asts.prog}: error: {e.args[0] if len(e.args) == 1 else e.args}\n```", TextParseMode.MarkDown)  # noqa: E501
         except Exception:
             re = await lib.editMessageText(mes['chat_id'], mes['id'], f"```\n{format_exc()}\n```", TextParseMode.MarkDown)  # noqa: E501
-    if re is None:
-        print('Can not edit message.')
+    if re is None or re is False:
+        print('Can not edit/delete message.')
 
 
 async def main(lib: TdLib, robot: TdLib):