refactor: 把 tempfile 修改成了 io.BytesIO

Author: tianxiu2b2t

.gitignore

@@ -7,6 +7,7 @@ config
 cache
 database/*.db
 cluster_status
+tianxiu2b2t
 
 **/__pycache__
 

core/cluster.py

@@ -1,6 +1,7 @@
 import contextlib
 import datetime
 import hmac
+import io
 import json
 import os
 from pathlib import Path
@@ -577,44 +578,40 @@ async def _download_file(
         for _ in range(10):
             size = 0
             hash = utils.get_hash_obj(file.hash)
-            with tempfile.NamedTemporaryFile(
-                dir=self._cache_dir,
-                delete=False
-            ) as tmp_file:
-                try:
-                    async with session.get(
-                        file.path
-                    ) as resp:
-                        while (data := await resp.content.read(1024 * 1024 * 16)):
-                            tmp_file.write(data)
-                            hash.update(data)
-                            inc = len(data)
-                            size += inc
-                            self._pbar.update(inc)
-                            pbar.update(inc)
-                        if hash.hexdigest() != file.hash or size != file.size:
-                            await anyio.sleep(50)
-                            raise Exception(f"hash mismatch, got {hash.hexdigest()} expected {file.hash}")
-                    await self.upload_storage(file, tmp_file, size)
-                    self.update_success()
-                except Exception as e:
-                    last_error = e
-                    self._pbar.update(-size)
-                    pbar.update(-size)
-                    self.update_failed()
-                    continue
-                finally:
-                    tmp_file.close()
-                    os.remove(tmp_file.name)
-                return None
+            tmp_file = io.BytesIO()
+            try:
+                async with session.get(
+                    file.path
+                ) as resp:
+                    while (data := await resp.content.read(1024 * 1024 * 16)):
+                        tmp_file.write(data)
+                        hash.update(data)
+                        inc = len(data)
+                        size += inc
+                        self._pbar.update(inc)
+                        pbar.update(inc)
+                    if hash.hexdigest() != file.hash or size != file.size:
+                        await anyio.sleep(50)
+                        raise Exception(f"hash mismatch, got {hash.hexdigest()} expected {file.hash}")
+                await self.upload_storage(file, tmp_file, size)
+                self.update_success()
+            except Exception as e:
+                last_error = e
+                self._pbar.update(-size)
+                pbar.update(-size)
+                self.update_failed()
+                continue
+            finally:
+                tmp_file.close()
+            return None
         if last_error is not None:
             raise last_error
 
 
     async def upload_storage(
         self,
         file: BMCLAPIFile,
-        tmp_file: 'tempfile._TemporaryFileWrapper',
+        data: io.BytesIO,
         size: int
     ):
         missing_storage = [
@@ -623,11 +620,11 @@ async def upload_storage(
         if len(missing_storage) == 0:
             return
         for storage in missing_storage:
-            tmp_file.seek(0)
+            data.seek(0)
             retries = 0
             while 1:
                 try:
-                    await storage.storage.upload_download_file(f"{file.hash[:2]}/{file.hash}", tmp_file, size)
+                    await storage.storage.upload_download_file(f"{file.hash[:2]}/{file.hash}", data, size)
                     break
                 except:
                     if retries >= 10:

core/dashboard.py

@@ -101,8 +101,8 @@ async def setup(
     app: fastapi.FastAPI,
     task_group: anyio.abc.TaskGroup,
 ):
-    if not DEBUG:
-        return
+    #if not DEBUG:
+    #    return
 
     systeminfo.setup()
 

core/logger.py

@@ -114,13 +114,16 @@ def emit(self, record):
 
 # 配置拦截处理器
 #logging.basicConfig(handlers=[InterceptHandler()], level=logging.DEBUG)
-#logging.getLogger("uvicorn").handlers = [InterceptHandler()]
-#logging.getLogger("uvicorn.access").handlers = [InterceptHandler()]
-#logging.getLogger("uvicorn.error").handlers = [InterceptHandler()]
-#logging.getLogger("engineio.client").handlers = [DebugHandler()]
-#logging.getLogger("socketio.client").handlers = [DebugHandler()]
+logging.getLogger("uvicorn").handlers = [InterceptHandler()]
+logging.getLogger("uvicorn.access").handlers = [InterceptHandler()]
+logging.getLogger("uvicorn.error").handlers = [InterceptHandler()]
+logging.getLogger("engineio.client").handlers.append(InterceptHandler())
+logging.getLogger("socketio.client").handlers.append(InterceptHandler())
 logging.basicConfig(
-    filename="./logs/logging.log"
+    #format=LOGGER_FORMAT,
+    handlers=[InterceptHandler()],
+    #filename="./logs/logging.log",
+    level=logging.DEBUG
 )
 
 

core/storage/__init__.py

@@ -21,7 +21,7 @@
     "local": LocalStorage,
     "alist": AlistStorage,
     "webdav": WebDavStorage,
-    "s3": S3Storage,
+    #"s3": S3Storage,
     "minio": MinioStorage
 }
 

core/storage/abc.py

@@ -2,6 +2,7 @@
 import tempfile
 
 import anyio.abc
+import io
 
 from core import utils
 from core.abc import BMCLAPIFile, ResponseFile, ResponseFileNotFound, ResponseFileMemory, ResponseFileLocal, ResponseFileRemote
@@ -151,15 +152,15 @@ async def works(root_ids: list[int]):
     async def upload(
         self,
         path: str,
-        tmp_file: tempfile._TemporaryFileWrapper,
+        data: io.BytesIO,
         size: int
     ):
         raise NotImplementedError
 
-    async def upload_download_file(self, path: str, tmp_file: tempfile._TemporaryFileWrapper, size: int):
+    async def upload_download_file(self, path: str, data: io.BytesIO, size: int):
         if self.download_dir:
             path = f"download/{path}"
-        await self.upload(f"download/{path}", tmp_file, size)
+        await self.upload(f"download/{path}", data, size)
 
     async def get_response_file(
         self,
@@ -197,15 +198,12 @@ async def write_measure(self, size: int):
         path = f"measure/{size}"
         size = size * 1024 * 1024
 
-        with tempfile.NamedTemporaryFile() as tmp:
-            tmp.write(b'\x00' * size)
-            tmp.seek(0)
-            await self.upload(
-                path,
-                tmp,
-                size
-            )
-            logger.tsuccess("storage.write_measure", size=int(size / (1024 * 1024)), name=self.name, type=self.type)
+        await self.upload(
+            path,
+            io.BytesIO(b"\x00" * size),
+            size
+        )
+        logger.tsuccess("storage.write_measure", size=int(size / (1024 * 1024)), name=self.name, type=self.type)
 
 
     def get_py_check_path(self) -> 'CPath':

core/storage/alist.py

@@ -1,3 +1,4 @@
+import io
 from tempfile import _TemporaryFileWrapper
 import time
 from typing import Any
@@ -167,7 +168,7 @@ async def list_files(self, path: str) -> list[abc.FileInfo]:
                     ))
         return res
 
-    async def upload(self, path: str, tmp_file: _TemporaryFileWrapper, size: int):
+    async def upload(self, path: str, data: io.BytesIO, size: int):
         async with aiohttp.ClientSession(
             base_url=self._endpoint,
             headers={
@@ -180,10 +181,10 @@ async def upload(self, path: str, tmp_file: _TemporaryFileWrapper, size: int):
                 headers={
                     "File-Path": urlparse.quote(str(self._path / path)),
                 },
-                data=tmp_file.file
+                data=data.getbuffer()
             ) as resp:
-                data = AlistResponse(await resp.json())
-                data.raise_for_status()
+                alist_resp = AlistResponse(await resp.json())
+                alist_resp.raise_for_status()
                 return True
 
     async def get_file(self, path: str) -> abc.ResponseFile:

core/storage/local.py

@@ -1,6 +1,5 @@
-import os
 from pathlib import Path
-import tempfile
+import io
 import time
 import anyio.abc
 
@@ -54,16 +53,13 @@ async def list_files(
     async def upload(
         self,
         path: str,
-        tmp_file: tempfile._TemporaryFileWrapper,
+        data: io.BytesIO,
         size: int
     ):
         root = Path(str(self.path)) / path
         root.parent.mkdir(parents=True, exist_ok=True)
         with open(root, "wb") as f:
-            while (data := tmp_file.read(65536)):
-                if not data:
-                    break
-                f.write(data)
+            f.write(data.getbuffer())
         return True
 
     async def _check(

core/storage/minio.py

@@ -94,15 +94,15 @@ async def list_files(
     async def upload(
         self,
         path: str,
-        tmp_file: tempfile._TemporaryFileWrapper,
+        data: io.BytesIO,
         size: int
     ):
         root = self.path / path
 
         await self.minio.put_object(
             self.bucket,
             str(root)[1:],
-            tmp_file,
+            data.getbuffer(),
             size
         )
         return True

core/storage/s3.py

@@ -1,4 +1,5 @@
 from io import BytesIO
+import io
 import tempfile
 import time
 import aioboto3.session
@@ -137,7 +138,7 @@ async def list_files(
     async def upload(
         self,
         path: str,
-        tmp_file: tempfile._TemporaryFileWrapper,
+        data: io.BytesIO,
         size: int
     ):
         async with self.session.resource(
@@ -149,7 +150,7 @@ async def upload(
         ) as resource:
             bucket = await resource.Bucket(self.bucket)
             obj = await bucket.Object(str(self.path / path))
-            await obj.upload_fileobj(tmp_file)
+            await obj.upload_fileobj(data)
         return True