From e82122c087d5e739352387838c611247935fd2a8 Mon Sep 17 00:00:00 2001
From: Akimio521 <akimio2333@gmail.com>
Date: Thu, 27 Jun 2024 20:24:10 +0800
Subject: [PATCH] =?UTF-8?q?feat(Alist2Strm):=E6=94=AF=E6=8C=81=E8=AE=BE?=
 =?UTF-8?q?=E7=BD=AE=20subtitle,img,nfo=20=E7=AD=89=E4=BA=8C=E8=BF=9B?=
 =?UTF-8?q?=E5=88=B6=E6=96=87=E4=BB=B6=E6=98=AF=E5=90=A6=E5=90=AF=E7=94=A8?=
 =?UTF-8?q?=E5=BC=82=E6=AD=A5=E4=B8=8B=E8=BD=BD=E6=88=96=E5=90=8C=E6=AD=A5?=
 =?UTF-8?q?=E4=B8=8B=E8=BD=BD?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 autofilm.py                |  1 +
 config/config.yaml.example |  3 +++
 modules/Alist2Strm.py      | 55 +++++++++++++++++++++++++-------------
 3 files changed, 41 insertions(+), 18 deletions(-)

diff --git a/autofilm.py b/autofilm.py
index ccc06f0..7aabae4 100644
--- a/autofilm.py
+++ b/autofilm.py
@@ -72,5 +72,6 @@ def run_Alist2Strm(self) -> None:
                     self.img,
                     self.nfo,
                     self.library_mode,
+                    alist_server.get("async_mode"),
                 )
                 alist2strm.run()
diff --git a/config/config.yaml.example b/config/config.yaml.example
index 365e26d..a43ddb2 100644
--- a/config/config.yaml.example
+++ b/config/config.yaml.example
@@ -12,15 +12,18 @@ AlistServerList:
     password: adminadmin             # Alist密码
     base_path: /ani/A/               # Alist服务器上文件夹路径 
     token:                           # AList未启用签名时，设置为空字符串
+    async_mode: False                # 设置 subtitle,img,nfo 等二进制文件是否启用异步下载
   - id: 电影A
     url: https://alist.example1.com
     username: alist
     password: alist
     base_path: /movie/A/      
     token:
+    async_mode: True
   - id: OneDrive
     url: https://alist.example2.com
     username: alist
     password: adminadmin
     base_path: /网盘/OneDrive/
     token: alist-a1b2c3d4-12...
+    async_mode: False
diff --git a/modules/Alist2Strm.py b/modules/Alist2Strm.py
index 3592529..8ea6e31 100644
--- a/modules/Alist2Strm.py
+++ b/modules/Alist2Strm.py
@@ -3,7 +3,8 @@
 import hashlib
 import base64
 import asyncio
-from aiohttp import ClientSession
+from requests import Session
+from aiohttp import ClientSession as AsyncSession
 from pathlib import Path
 from typing import Optional
 
@@ -28,6 +29,7 @@ def __init__(
         img: bool = False,
         nfo: bool = False,
         library_mode: bool = True,
+        async_mode: bool = False,
     ) -> None:
         self.alist_server_url = alist_server_url.rstrip("/")
         self.alist_server_username = alist_server_username
@@ -42,6 +44,7 @@ def __init__(
         self.img = img
         self.nfo = nfo
         self.library_mode = library_mode
+        self.async_mode = async_mode
 
         logging.debug(
             f"Alist2Strm配置".center(50, "=") + "\n"
@@ -54,7 +57,8 @@ def __init__(
             f"是否下载字幕：{self.subtitle}\n"
             f"是否下载图片：{self.img}\n"
             f"是否下载NFO：{self.nfo}\n"
-            f"是否为库模式：{self.library_mode}"
+            f"是否为库模式：{self.library_mode}\n"
+            f"是否为启用异步下载：{self.async_mode}"
         )
 
     def run(self) -> None:
@@ -87,16 +91,24 @@ async def _processer(self):
             )
             return
 
-        async with ClientSession() as session:
-            tasks = [
-                asyncio.create_task(self._file_processer(alist_path_cls, session))
-                for alist_path_cls in fs.rglob("*.*")
-            ]
-            await asyncio.gather(*tasks)
+        if self.async_mode:
+            self.session = AsyncSession()
+        else:
+            self.session = Session()
 
-    async def _file_processer(
-        self, alist_path_cls: AlistPath, session: ClientSession
-    ) -> None:
+        tasks = [
+            asyncio.create_task(self._file_processer(alist_path_cls))
+            for alist_path_cls in fs.rglob("*.*")
+        ]
+        await asyncio.gather(*tasks)
+
+        if self.session:
+            if self.async_mode:
+                await self.session.close()
+            else:
+                self.session.close()
+
+    async def _file_processer(self, alist_path_cls: AlistPath) -> None:
         if not alist_path_cls.name.lower().endswith(ALL_EXT):
             return
 
@@ -140,13 +152,20 @@ async def _file_processer(
             return
         else:
             file_output_path.parent.mkdir(parents=True, exist_ok=True)
-            async with session.get(file_download_url) as resp:
-                if resp.status == 200:
-                    with file_output_path.open(mode="wb") as f:
-                        f.write(await resp.read())
-                    logging.debug(
-                        f"{file_output_path.name}下载成功，文件本地目录：{file_output_path.parent}"
-                    )
+
+            if self.async_mode:
+                resp = await self.session.get(file_download_url)
+            else:
+                resp = self.session.get(file_download_url)
+
+            with file_output_path.open(mode="wb") as f:
+                if self.async_mode:
+                    f.write(await resp.read())
+                else:
+                    f.write(resp.content)
+            logging.debug(
+                f"{file_output_path.name}下载成功，文件本地目录：{file_output_path.parent}"
+            )
 
     def _sign(self, secret_key: Optional[str], data: str) -> str:
         if secret_key == "" or secret_key == None: