添加aria2的支持

foyoux · Sep 5, 2021 · c74a044 · c74a044
1 parent 20f6612
commit c74a044
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 8 deletions.
diff --git a/aligo/core/BaseAligo.py b/aligo/core/BaseAligo.py
@@ -1,5 +1,6 @@
 """..."""
 import json
+import subprocess
 import traceback
 from dataclasses import asdict
 from typing import Generic, List, Iterator, Dict, Callable
@@ -19,13 +20,15 @@
 class BaseAligo:
     """..."""
 
-    def __init__(self, auth: Optional[Auth] = None):
+    def __init__(self, use_aria2: bool = True, auth: Optional[Auth] = None):
         self._auth: Auth = auth or Auth()
         self._session: requests.Session = self._auth.session
         self._token: Token = self._auth.token
         self._user: Optional[BaseUser] = None
         self._personal_info: Optional[GetPersonalInfoResponse] = None
         self._default_drive: Optional[BaseDrive] = None
+        self._has_aria2c = True if subprocess.run(['aria2c', '-h'],
+                                                  capture_output=True).returncode == 0 and use_aria2 else False
 
     def _post(self, path: str, host: str = API_HOST, body: Union[DataType, Dict] = None) -> requests.Response:
         """统一处理数据类型和 drive_id"""

diff --git a/aligo/core/Download.py b/aligo/core/Download.py
@@ -49,20 +49,33 @@ def download_file(self, file_path: str, url: str) -> str:
         :param url: 下载地址
         :return: 下载完成保存文件的本地路径
         """
-        file_dir, file_name = os.path.split(file_path)
+        file_dir, file_name = os.path.split(os.path.abspath(file_path))
         file_name = self._del_special_symbol(file_name)
-        file_path = os.path.abspath(os.path.join(file_dir, file_name))
+        file_path = os.path.join(file_dir, file_name)
 
-        # 递归创建目录
-        os.makedirs(os.path.dirname(file_path), exist_ok=True)
+        self._auth.log.info(f'开始下载文件 {file_path}')
 
-        if os.path.exists(file_path):
+        if os.path.exists(file_path) and not os.path.exists(f'{file_path}.aria2'):
             self._auth.log.warning(f'文件已存在,跳过下载 {file_path}')
             return file_path
 
-        tmp_file = file_path + '.ali'
+        if self._has_aria2c:
+            cmd = ' '.join([
+                f'aria2c "{url}"',
+                f'--referer=https://www.aliyundrive.com/',
+                f'--console-log-level=warn',
+                f'--download-result=hide',
+                f'--auto-file-renaming=false',
+                f'--force-sequential=true',
+                f'-o "{file_path}"',
+            ])
+            # print(cmd)
+            os.system(cmd)
+            return file_path
 
-        self._auth.log.info(f'开始下载文件 {file_path}')
+        # 递归创建目录
+        os.makedirs(os.path.dirname(file_path), exist_ok=True)
+        tmp_file = file_path + '.ali'
 
         with requests.get(url, headers={
             'referer': 'https://www.aliyundrive.com/'