修复下载功能的问题，待完善tab触发功能

DrissionPage/browser_download_manager.py

@@ -1,4 +1,5 @@
 # -*- coding:utf-8 -*-
+from os.path import sep
 from pathlib import Path
 from shutil import move
 from time import sleep, perf_counter
@@ -117,31 +118,41 @@ class BrowserDownloadManager(object):
         if mission.state not in ('canceled', 'skipped'):
             mission.state = state
         mission.final_path = final_path
-        if mission.tab_id in self._tab_missions:
+        if mission.tab_id in self._tab_missions and mission.id in self._tab_missions[mission.tab_id]:
             self._tab_missions[mission.tab_id].remove(mission.id)
         self._missions.pop(mission.id)
 
-    def cancel(self, mission, state):
+    def cancel(self, mission):
         """取消任务
         :param mission: 任务对象
-        :param state: 任务状态
         :return: None
         """
-        mission.state = state
+        mission.state = 'canceled'
         self._page.browser_driver.call_method('Browser.cancelDownload', guid=mission.id)
         if mission.final_path:
             Path(mission.final_path).unlink(True)
 
+    def skip(self, mission):
+        """跳过任务
+        :param mission: 任务对象
+        :return: None
+        """
+        mission.state = 'skipped'
+        self._page.browser_driver.call_method('Browser.cancelDownload', guid=mission.id)
+
     def _onDownloadWillBegin(self, **kwargs):
         """用于获取弹出新标签页触发的下载任务"""
         guid = kwargs['guid']
-        end = perf_counter() + .3
+
+        end = perf_counter() + 2
         while perf_counter() < end:
             tab_id = self._guid_and_tab.get(guid, None)
             if tab_id:
+                # print('拿到')
                 break
             sleep(.005)
         else:
+            # print('没拿到')
             tab_id = self._page.tab_id
 
         settings = TabDownloadSettings(tab_id)
@@ -164,13 +175,13 @@ class BrowserDownloadManager(object):
             elif settings.when_file_exists == 'overwrite':
                 goal_path.unlink()
 
-        m = DownloadMission(self, tab_id, guid, settings.path, name, kwargs['url'])
+        m = DownloadMission(self, tab_id, guid, settings.path, name, kwargs['url'], self._page.download_path)
         self._missions[guid] = m
 
         if self.get_flag(tab_id) is False:  # 取消该任务
-            self.cancel(m, 'canceled')
+            self.cancel(m)
         elif skip:
-            self.cancel(m, 'skipped')
+            self.skip(m)
         else:
             self._tab_missions.setdefault(tab_id, []).append(guid)
 
@@ -181,15 +192,18 @@ class BrowserDownloadManager(object):
         if kwargs['guid'] in self._missions:
             mission = self._missions[kwargs['guid']]
             if kwargs['state'] == 'inProgress':
-                mission.state = 'running'
                 mission.received_bytes = kwargs['receivedBytes']
                 mission.total_bytes = kwargs['totalBytes']
 
             elif kwargs['state'] == 'completed':
+                if mission.state == 'skipped':
+                    Path(f'{mission.save_path}{sep}{mission.id}').unlink(True)
+                    self.set_done(mission, 'skipped')
+                    return
                 mission.received_bytes = kwargs['receivedBytes']
                 mission.total_bytes = kwargs['totalBytes']
-                form_path = f'{self._page.download_path}\\{mission.id}'
+                form_path = f'{mission.save_path}{sep}{mission.id}'
-                to_path = str(get_usable_path(f'{mission.path}\\{mission.name}'))
+                to_path = str(get_usable_path(f'{mission.path}{sep}{mission.name}'))
                 move(form_path, to_path)
                 self.set_done(mission, 'completed', final_path=to_path)
 
@@ -224,17 +238,27 @@ class TabDownloadSettings(object):
 
 
 class DownloadMission(object):
-    def __init__(self, mgr, tab_id, _id, path, name, url):
+    def __init__(self, mgr, tab_id, _id, path, name, url, save_path):
+        """
+        :param mgr: BrowserDownloadManager对象
+        :param tab_id: 标签页id
+        :param _id: 任务id
+        :param path: 保存路径
+        :param name: 文件名
+        :param url: url
+        :param save_path: 下载路径
+        """
         self._mgr = mgr
         self.url = url
         self.tab_id = tab_id
         self.id = _id
         self.path = path
         self.name = name
-        self.state = 'waiting'
+        self.state = 'running'
         self.total_bytes = None
         self.received_bytes = 0
         self.final_path = None
+        self.save_path = save_path
 
     def __repr__(self):
         return f'<DownloadMission {id(self)} {self.rate}>'
@@ -247,11 +271,11 @@ class DownloadMission(object):
     @property
     def is_done(self):
         """返回任务是否在运行中"""
-        return self.state == 'completed'
+        return self.state in ('completed', 'skipped', 'canceled')
 
     def cancel(self):
         """取消该任务，如任务已完成，删除已下载的文件"""
-        self._mgr.cancel(self, state='canceled')
+        self._mgr.cancel(self)
 
     def wait(self, show=True, timeout=None, cancel_if_timeout=True):
         """等待任务结束
@@ -269,23 +293,19 @@ class DownloadMission(object):
             print(f'目标路径：{self.path}')
 
         if timeout is None:
-            while self.id in self._mgr.missions:
+            while not self.is_done:
                 if show:
                     print(f'\r{self.rate}% ', end='')
                 sleep(.2)
 
         else:
-            running = True
             end_time = perf_counter() + timeout
             while perf_counter() < end_time:
                 if show:
                     print(f'\r{self.rate}% ', end='')
-                if self.id not in self._mgr.missions:
-                    running = False
-                    break
                 sleep(.2)
 
-            if running and cancel_if_timeout:
+            if not self.is_done and cancel_if_timeout:
                 self.cancel()
 
         if show:

DrissionPage/browser_download_manager.pyi

@@ -36,7 +36,9 @@ class BrowserDownloadManager(object):
 
     def set_done(self, mission: DownloadMission, state: str, final_path: str = None) -> None: ...
 
-    def cancel(self, mission: DownloadMission, state: str) -> None: ...
+    def cancel(self, mission: DownloadMission) -> None: ...
+
+    def skip(self, mission: DownloadMission) -> None: ...
 
     def _onDownloadWillBegin(self, **kwargs) -> None: ...
 
@@ -65,8 +67,10 @@ class DownloadMission(object):
     total_bytes: Optional[int] = ...
     received_bytes: int = ...
     final_path: Optional[str] = ...
+    save_path: str = ...
 
-    def __init__(self, mgr: BrowserDownloadManager, tab_id: str, _id: str, path: str, name: str, url: str): ...
+    def __init__(self, mgr: BrowserDownloadManager, tab_id: str, _id: str, path: str, name: str, url: str,
+                 save_path: str): ...
 
     @property
     def rate(self) -> float: ...

DrissionPage/chromium_base.py

@@ -250,6 +250,7 @@ class ChromiumBase(BasePage):
 
     def _onDownloadWillBegin(self, **kwargs):
         """下载即将开始时执行"""
+        print('aaa')
         self.browser._dl_mgr.set_mission(self.tab_id, kwargs['guid'])
 
     def __call__(self, loc_or_str, timeout=None):

DrissionPage/chromium_page.pyi

@@ -66,7 +66,7 @@ class ChromiumPage(ChromiumBase):
     def get_tab(self, tab_id: Union[str, ChromiumTab] = None) -> ChromiumTab: ...
 
     def find_tabs(self, title: str = None, url: str = None,
-                  tab_type: Union[str, list, tuple, set] = None, single: bool = True) -> Union[str, List[str]]: ...
+                  tab_type: Union[str, list, tuple] = None, single: bool = True) -> Union[str, List[str]]: ...
 
     def _new_tab(self, url: str = None, switch_to: bool = False) -> str: ...
 

DrissionPage/waiter.py

@@ -250,7 +250,7 @@ class ChromiumPageWaiter(ChromiumTabWaiter):
 
             if self._driver._dl_mgr._missions:
                 if cancel_if_timeout:
-                    for m in self._driver._dl_mgr._missions.values():
+                    for m in list(self._driver._dl_mgr._missions.values()):
                         m.cancel()
                 return False
             else: