From 9a190484bfafcf43f208fa5ff3b8f2921055369c Mon Sep 17 00:00:00 2001
From: veoco <one@nomox.cn>
Date: Tue, 13 Dec 2022 14:18:02 +0800
Subject: [PATCH] =?UTF-8?q?=E8=B0=83=E6=95=B4=E6=96=87=E4=BB=B6=E5=88=A0?=
 =?UTF-8?q?=E9=99=A4=E6=96=B9=E5=BC=8F?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 main.py    | 19 +++++++++++++------
 storage.py | 15 ++++++---------
 2 files changed, 19 insertions(+), 15 deletions(-)

diff --git a/main.py b/main.py
index be822f6..5ba5525 100644
--- a/main.py
+++ b/main.py
@@ -51,9 +51,13 @@ async def delete_expire_files():
         async with AsyncSession(engine, expire_on_commit=False) as s:
             query = select(Codes).where(or_(Codes.exp_time < datetime.datetime.now(), Codes.count == 0))
             exps = (await s.execute(query)).scalars().all()
-            files = [{'type': old.type, 'text': old.text} for old in exps]
+            files = []
+            exps_ids = []
+            for exp in exps:
+                if exp.type != "text":
+                    files.append(exp.text)
+                exps_ids.append(exp.id)
             await storage.delete_files(files)
-            exps_ids = [exp.id for exp in exps]
             query = delete(Codes).where(Codes.id.in_(exps_ids))
             await s.execute(query)
             await s.commit()
@@ -83,9 +87,11 @@ async def admin_post(s: AsyncSession = Depends(get_session)):
 async def admin_delete(code: str, s: AsyncSession = Depends(get_session)):
     query = select(Codes).where(Codes.code == code)
     file = (await s.execute(query)).scalars().first()
-    await storage.delete_file({'type': file.type, 'text': file.text})
-    await s.delete(file)
-    await s.commit()
+    if file:
+        if file.type != 'text':
+            await storage.delete_file(file.text)
+        await s.delete(file)
+        await s.commit()
     return {'detail': '删除成功'}
 
 
@@ -115,7 +121,8 @@ async def index(code: str, ip: str = Depends(ip_limit), s: AsyncSession = Depend
         error_count = settings.ERROR_COUNT - ip_limit.add_ip(ip)
         raise HTTPException(status_code=404, detail=f"取件码错误，错误{error_count}次将被禁止10分钟")
     if info.exp_time < datetime.datetime.now() or info.count == 0:
-        await storage.delete_file({'type': info.type, 'text': info.text})
+        if info.type != "text":
+            await storage.delete_file(info.text)
         await s.delete(info)
         await s.commit()
         raise HTTPException(status_code=404, detail="取件码已过期，请联系寄件人")
diff --git a/storage.py b/storage.py
index 58a90de..b992a9a 100644
--- a/storage.py
+++ b/storage.py
@@ -46,16 +46,13 @@ class FileSystemStorage:
         filepath = await self.get_filepath(text)
         await asyncio.to_thread(self._save, filepath, file.file)
 
-    async def delete_file(self, file):
-        # 是文件就删除
-        if file['type'] != 'text':
-            filepath = self.DATA_ROOT / file['text'].lstrip(self.STATIC_URL + '/')
-            await asyncio.to_thread(os.remove, filepath)
+    async def delete_file(self, text: str):
+        filepath = await self.get_filepath(text)
+        await asyncio.to_thread(os.remove, filepath)
 
-    async def delete_files(self, files):
-        for file in files:
-            if file['type'] != 'text':
-                await self.delete_file(file)
+    async def delete_files(self, texts):
+        tasks = [self.delete_file(text) for text in texts]
+        await asyncio.gather(*tasks)
 
 
 STORAGE_ENGINE = {