[PATCH v3 13/17] daemon: Add @handle_flood decorator and remove some functions

[Muhammad Rizki <kiizuha@gnuweeb.org>]

From: Muhammad Rizki <riskimuhammmad1@gmail.com>

I decided to create @handle_flood decorator to make it easier to handle
the Telegram floodwait. Remove __send_text_msg() and replace the
call of the __send_text_msg() with send_text_email() in
packages/client.py and it have handle_flood decorator support.

Signed-off-by: Muhammad Rizki <riskimuhammmad1@gmail.com>
---
 daemon/packages/client.py    |  3 +++
 daemon/packages/decorator.py | 38 +++++++++++++++++++++++++++++
 daemon/scraper/bot.py        | 47 ++++++------------------------------
 3 files changed, 49 insertions(+), 39 deletions(-)
 create mode 100644 daemon/packages/decorator.py

diff --git a/daemon/packages/client.py b/daemon/packages/client.py
index f73b913..45acac3 100644
--- a/daemon/packages/client.py
+++ b/daemon/packages/client.py
@@ -7,6 +7,7 @@ from pyrogram import Client
 from pyrogram.enums import ParseMode
 from pyrogram.types import Message, InlineKeyboardMarkup, InlineKeyboardButton
 from typing import Union, BinaryIO
+from .decorator import handle_flood
 
 
 class DaemonClient(Client):
@@ -16,6 +17,7 @@ class DaemonClient(Client):
 				api_hash, **kwargs)
 
 
+	@handle_flood
 	async def send_text_email(
 		self,
 		chat_id: Union[int, str],
@@ -39,6 +41,7 @@ class DaemonClient(Client):
 		)
 
 
+	@handle_flood
 	async def send_patch_email(
 		self,
 		chat_id: Union[int, str],
diff --git a/daemon/packages/decorator.py b/daemon/packages/decorator.py
new file mode 100644
index 0000000..7d7dc39
--- /dev/null
+++ b/daemon/packages/decorator.py
@@ -0,0 +1,38 @@
+# SPDX-License-Identifier: GPL-2.0-only
+#
+# Copyright (C) 2022  Muhammad Rizki <riskimuhammmad1@gmail.com>
+#
+
+from pyrogram.errors.exceptions.flood_420 import FloodWait
+from typing import Any, Callable
+from functools import wraps
+import re
+import asyncio
+
+__all__ = ["handle_flood"]
+
+
+def handle_flood(func: Callable[[Any], Any]) -> Callable[[Any], Any]:
+	@wraps(func)
+	async def callback(*args: Any) -> Any:
+		while True:
+			try:
+				return await func(*args)
+			except FloodWait as e:
+				#
+				# Aiee... we hit our limit.
+				# Let's slow down a bit.
+				#
+				_flood_exceptions(e)
+				print("[__handle_telegram_floodwait]: Woken up from flood wait...")
+	return callback
+
+
+async def _flood_exceptions(e):
+	x = re.search(r"A wait of (\d+) seconds is required", str(e))
+	if not x:
+		raise e
+
+	n = int(x.group(1))
+	print(f"[____handle_telegram_floodwait]: Sleeping for {n} seconds due to Telegram limit")
+	await asyncio.sleep(n)
diff --git a/daemon/scraper/bot.py b/daemon/scraper/bot.py
index 1abfa98..38c6e94 100644
--- a/daemon/scraper/bot.py
+++ b/daemon/scraper/bot.py
@@ -37,7 +37,7 @@ class Bot():
 
 
 	def __init__(self, client: DaemonClient, sched: AsyncIOScheduler,
-		     scraper: Scraper, mutexes: BotMutexes, conn):
+			scraper: Scraper, mutexes: BotMutexes, conn):
 		self.client = client
 		self.sched = sched
 		self.scraper = scraper
@@ -88,7 +88,7 @@ class Bot():
 		for tg_chat_id in self.TG_CHAT_IDS:
 			async with self.mutexes.send_to_tg:
 				should_wait = await self.__send_to_tg(url, mail,
-								      tg_chat_id)
+									tg_chat_id)
 
 			if should_wait:
 				await asyncio.sleep(1)
@@ -105,7 +105,7 @@ class Bot():
 			return False
 
 		email_id = self.__need_to_send_to_telegram(email_msg_id,
-							   tg_chat_id)
+							tg_chat_id)
 		if not email_id:
 			#
 			# Email has already been sent to Telegram.
@@ -122,8 +122,9 @@ class Bot():
 							reply_to, text, url)
 		else:
 			text = "#ml\n" + text
-			m = await self.__send_text_msg(tg_chat_id, text,
-						       reply_to, url)
+			m = await self.client.send_text_email(
+				tg_chat_id, text,reply_to, url
+			)
 
 		self.db.insert_telegram(email_id, m.chat.id, m.id)
 		for d, f in files:
@@ -161,40 +162,8 @@ class Bot():
 		print("[__send_patch_msg]")
 
 		tmp, doc, caption, url = utils.prepare_send_patch(mail, text, url)
-		ret = await self.__handle_telegram_floodwait(
-			self.client.send_patch_email,
-			*[tg_chat_id, doc, caption, reply_to, url]
+		ret = await self.client.send_patch_email(
+			tg_chat_id, doc, caption, reply_to, url
 		)
 		utils.clean_up_after_send_patch(tmp)
 		return ret
-
-
-	async def __send_text_msg(self, *args):
-		return await self.__handle_telegram_floodwait(
-			self.client.send_text_email,
-			*args
-		)
-
-
-	async def __handle_telegram_floodwait(self, callback, *args):
-		while True:
-			try:
-				return await callback(*args)
-			except pyrogram.errors.exceptions.flood_420.FloodWait as e:
-				#
-				# Aiee... we hit our limit.
-				# Let's slow down a bit.
-				#
-				await self.____handle_telegram_floodwait(e)
-				print("[__handle_telegram_floodwait]: Woken up from flood wait...")
-
-
-	async def ____handle_telegram_floodwait(self, e):
-		x = str(e)
-		x = re.search(r"A wait of (\d+) seconds is required", x)
-		if not x:
-			raise e
-
-		n = int(x.group(1))
-		print(f"[____handle_telegram_floodwait]: Sleeping for {n} seconds due to Telegram limit")
-		await asyncio.sleep(n)
-- 
Muhammad Rizki