调整文件大小计算和写入方法

storage.py

@@ -1,7 +1,10 @@
 import os
 import asyncio
-import datetime
+from datetime import datetime
 from pathlib import Path
+from typing import BinaryIO
+
+from fastapi import UploadFile
 
 import settings
 
@@ -11,25 +14,38 @@ class FileSystemStorage:
     STATIC_URL = settings.STATIC_URL
     NAME = "filesystem"
 
-    async def get_filepath(self, path):
-        return self.DATA_ROOT / path.lstrip(self.STATIC_URL + '/')
+    async def get_filepath(self, text: str):
+        return self.DATA_ROOT / text.lstrip(self.STATIC_URL + '/')
 
-    def _save(self, filepath, file_bytes):
-        with open(filepath, 'wb') as f:
-            f.write(file_bytes)
-
-    async def save_file(self, file, file_bytes, key):
-        now = datetime.datetime.now()
-        path = self.DATA_ROOT / f"upload/{now.year}/{now.month}/{now.day}/"
+    async def get_text(self, file: UploadFile, key: str):
         ext = file.filename.split('.')[-1]
-        name = f'{key}.{ext}'
+        now = datetime.now()
+        path = self.DATA_ROOT / f"upload/{now.year}/{now.month}/{now.day}/"
         if not path.exists():
             path.mkdir(parents=True)
-        filepath = path / name
-        await asyncio.to_thread(self._save, filepath, file_bytes)
+        filepath = path / f'{key}.{ext}'
         text = f"{self.STATIC_URL}/{filepath.relative_to(self.DATA_ROOT)}"
         return text
 
+    async def get_size(self, file: UploadFile):
+        f = file.file
+        f.seek(0, os.SEEK_END)
+        size = f.tell()
+        f.seek(0, os.SEEK_SET)
+        return size
+
+    def _save(self, filepath, file: BinaryIO):
+        with open(filepath, 'wb') as f:
+            chunk_size = 256 * 1024
+            chunk = file.read(chunk_size)
+            while chunk:
+                f.write(chunk)
+                chunk = file.read(chunk_size)
+
+    async def save_file(self, file: UploadFile, text: str):
+        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':