[PATCH v4 4/9] atom: Improve remove_patch()

[Muhammad Rizki <kiizuha@gnuweeb.org>]

Improvement on remove_patch(). So, after send all the email attachments
then remove them. On the previous patch, this function can't remove
them properly.

Signed-off-by: Muhammad Rizki <kiizuha@gnuweeb.org>
---
 daemon/atom/utils.py                                | 10 +++++++---
 daemon/dscord/mailer/listener.py                    |  4 ++--
 daemon/telegram/mailer/listener.py                  |  3 +--
 daemon/telegram/packages/plugins/commands/scrape.py |  3 +--
 4 files changed, 11 insertions(+), 9 deletions(-)

diff --git a/daemon/atom/utils.py b/daemon/atom/utils.py
index a30d5cb..f554f6f 100644
--- a/daemon/atom/utils.py
+++ b/daemon/atom/utils.py
@@ -6,7 +6,7 @@
 
 from pyrogram.types import Chat, InlineKeyboardMarkup, InlineKeyboardButton
 from email.message import Message
-from typing import Dict
+from typing import Dict, Union
 from slugify import slugify
 import hashlib
 import uuid
@@ -238,8 +238,12 @@ def prepare_patch(mail, text, url, platform: str):
 	return tmp, file, caption, url
 
 
-def remove_patch(tmp):
-	shutil.rmtree(tmp)
+def remove_patch(tmp: Union[str, list]):
+	if isinstance(tmp, str):
+		return shutil.rmtree(tmp)
+
+	for d,_ in tmp:
+		shutil.rmtree(d)
 
 
 def fix_utf8_char(text: str, html_escape: bool = True):
diff --git a/daemon/dscord/mailer/listener.py b/daemon/dscord/mailer/listener.py
index a280a58..cc0a9f7 100644
--- a/daemon/dscord/mailer/listener.py
+++ b/daemon/dscord/mailer/listener.py
@@ -125,10 +125,10 @@ class Listener:
 
 		for d, f in files:
 			await m.reply(file=File(f"{d}/{f}"))
-			if files.index((d,f)) == len(files)-1:
-				utils.remove_patch(d)
 			await asyncio.sleep(1)
 
+		utils.remove_patch(files)
+
 		return True
 
 
diff --git a/daemon/telegram/mailer/listener.py b/daemon/telegram/mailer/listener.py
index 208aed0..08feddd 100644
--- a/daemon/telegram/mailer/listener.py
+++ b/daemon/telegram/mailer/listener.py
@@ -118,8 +118,7 @@ class Bot():
 			await m.reply_document(f"{d}/{f}", file_name=f)
 			await asyncio.sleep(1)
 
-		if files:
-			shutil.rmtree(str(files[0][0]))
+		utils.remove_patch(files)
 
 		return True
 
diff --git a/daemon/telegram/packages/plugins/commands/scrape.py b/daemon/telegram/packages/plugins/commands/scrape.py
index d4d10a9..52ddb0b 100644
--- a/daemon/telegram/packages/plugins/commands/scrape.py
+++ b/daemon/telegram/packages/plugins/commands/scrape.py
@@ -53,5 +53,4 @@ async def scrap_email(c: DaemonClient, m: Message):
 		await m.reply_document(f"{d}/{f}", file_name=f)
 		await asyncio.sleep(1)
 
-	if files:
-		shutil.rmtree(str(files[0][0]))
+	utils.remove_patch(files)
-- 
Muhammad Rizki