改进下载功能基本完成

DrissionPage/chromium_driver.py

@@ -79,7 +79,7 @@ class ChromiumDriver(object):
         message_json = dumps(message)
 
         if self.debug:
-            if isinstance(self.debug, str) and message.get('method', '').startswith(self.debug):
+            if self.debug is True or (isinstance(self.debug, str) and message.get('method', '').startswith(self.debug)):
                 print(f'发> {message_json}')
             elif isinstance(self.debug, (list, tuple, set)):
                 for m in self.debug:
@@ -131,7 +131,8 @@ class ChromiumDriver(object):
                 return
 
             if self.debug:
-                if 'id' in mes or (isinstance(self.debug, str) and mes.get('method', '').startswith(self.debug)):
+                if self.debug is True or 'id' in mes or (isinstance(self.debug, str)
+                                                         and mes.get('method', '').startswith(self.debug)):
                     print(f'<收 {message_json}')
                 elif isinstance(self.debug, (list, tuple, set)):
                     for m in self.debug:

DrissionPage/chromium_page.py

@@ -383,13 +383,20 @@ class ChromiumPageWaiter(ChromiumBaseWaiter):
         super().__init__(page)
         self._listener = None
 
-    def download_begin(self, timeout=None):
+    def download_begin(self, timeout=1.5):
         """等待浏览器下载开始
         :param timeout: 等待超时时间，为None则使用页面对象timeout属性
         :return: 是否等到下载开始
         """
         return self._driver.download_set.wait_download_begin(timeout)
 
+    def download_finish(self, timeout=None):
+        """等待下载结束
+        :param timeout: 等待超时时间，为None则使用页面对象timeout属性
+        :return: 是否等到下载结束
+        """
+        return self._driver.download_set.wait_download_finish(timeout)
+
     def new_tab(self, timeout=None):
         """等待新标签页出现
         :param timeout: 等待超时时间，为None则使用页面对象timeout属性
@@ -487,6 +494,7 @@ class ChromiumDownloadSetter(DownloadSetter):
         self._waiting_download = False
         self._download_begin = False
         self._browser_missions = {}
+        self._browser_downloading_count = 0
 
     @property
     def session(self):
@@ -495,6 +503,16 @@ class ChromiumDownloadSetter(DownloadSetter):
             self._session = Session()
         return self._session
 
+    @property
+    def browser_missions(self):
+        """返回浏览器下载任务"""
+        return list(self._browser_missions.values())
+
+    @property
+    def DownloadKit_missions(self):
+        """返回DownloadKit下载任务"""
+        return list(self.DownloadKit.missions.values())
+
     @property
     def _switched_DownloadKit(self):
         """返回从浏览器同步cookies后的Session对象"""
@@ -540,7 +558,7 @@ class ChromiumDownloadSetter(DownloadSetter):
         try:
             self._page.browser_driver.Browser.setDownloadBehavior(behavior='deny', eventsEnabled=True)
             self._page.browser_driver.Browser.downloadWillBegin = self._download_by_DownloadKit
-            self._page.browser_driver.Browser.downloadProgress = None
+            # self._page.browser_driver.Browser.downloadProgress = None
         except CallMethodError:
             raise RuntimeError('您的浏览器版本太低，不支持此方法，请升级。')
 
@@ -564,6 +582,19 @@ class ChromiumDownloadSetter(DownloadSetter):
         self._waiting_download = False
         return result
 
+    def wait_download_finish(self, timeout=None):
+        """等待所有下载结束
+        :param timeout: 超时时间
+        :return: 是否等待到下载完成
+        """
+        timeout = timeout if timeout is not None else self._page.timeout
+        end_time = perf_counter() + timeout
+        while perf_counter() < end_time:
+            if (self._DownloadKit is None or not self.DownloadKit.is_running) and self._browser_downloading_count == 0:
+                return True
+            sleep(.5)
+        return False
+
     def _cookies_to_session(self):
         """把driver对象的cookies复制到session对象"""
         ua = self._page.run_cdp('Runtime.evaluate', expression='navigator.userAgent;')['result']['value']
@@ -589,6 +620,7 @@ class ChromiumDownloadSetter(DownloadSetter):
 
         if self._waiting_download:
             self._download_begin = True
+        self._browser_downloading_count += 1
 
     def _download_will_begin(self, **kwargs):
         """浏览器下载即将开始时调用"""
@@ -604,24 +636,30 @@ class ChromiumDownloadSetter(DownloadSetter):
 
         if self._waiting_download:
             self._download_begin = True
+        self._browser_downloading_count += 1
 
     def _download_progress(self, **kwargs):
         """下载状态产生变化时调用"""
         guid = kwargs['guid']
-        m = self._browser_missions[guid]
+        m = self._browser_missions.get(guid, None)
-        m.size = kwargs['totalBytes']
+        if m:
-        m.received = kwargs['receivedBytes']
+            m.size = kwargs['totalBytes']
-        m.state = kwargs['state']
+            m.received = kwargs['receivedBytes']
+            m.state = kwargs['state']
 
-        if m.state == 'completed':
+            if m.state == 'completed':
-            path = Path(self._save_path) / m.name
+                path = Path(self._save_path) / m.name
-            from_path = Path(self._save_path) / guid
+                from_path = Path(self._save_path) / guid
-            if path.exists():
+                if path.exists():
-                if self._file_exists == 'rename':
+                    if self._file_exists == 'rename':
-                    path = get_usable_path(path)
+                        path = get_usable_path(path)
-                else:  # 'overwrite'
+                    else:  # 'overwrite'
-                    path.unlink()
+                        path.unlink()
-            from_path.rename(path)
+                from_path.rename(path)
+                m.save_path = path.absolute()
+
+        if kwargs['state'] != 'inProgress':
+            self._browser_downloading_count -= 1
 
     def _wait_download_complete(self):
         """等待DownloadKit下载完成"""
@@ -638,6 +676,9 @@ class BrowserDownloadMission(object):
         self.size = None
         self.received = None
 
+    def __repr__(self):
+        return f'<BrowserDownloadMission {self.save_path}>'
+
 
 class Alert(object):
     """用于保存alert信息的类"""

DrissionPage/chromium_page.pyi

@@ -9,6 +9,7 @@ from threading import Thread
 from typing import Union, Tuple, List, Dict
 
 from DownloadKit import DownloadKit
+from DownloadKit.mission import Mission
 from requests import Session
 
 from .chromium_base import ChromiumBase, ChromiumBaseSetter, ChromiumBaseWaiter, NetworkListener
@@ -113,7 +114,9 @@ class ChromiumPageWaiter(ChromiumBaseWaiter):
     _driver: ChromiumPage = ...
     _listener: Union[NetworkListener, None] = ...
 
-    def download_begin(self, timeout: float = None) -> bool: ...
+    def download_begin(self, timeout: float = 1.5) -> bool: ...
+
+    def download_finish(self, timeout: float = None) -> bool: ...
 
     def new_tab(self, timeout: float = None) -> bool: ...
 
@@ -161,10 +164,17 @@ class ChromiumDownloadSetter(DownloadSetter):
         self._waiting_download: bool = ...
         self._download_begin: bool = ...
         self._browser_missions: Dict[str, BrowserDownloadMission] = ...
+        self._browser_downloading_count: int = ...
 
     @property
     def session(self) -> Session: ...
 
+    @property
+    def browser_missions(self) -> List[BrowserDownloadMission]: ...
+
+    @property
+    def DownloadKit_missions(self) -> List[Mission]: ...
+
     @property
     def _switched_DownloadKit(self) -> DownloadKit: ...
 
@@ -176,6 +186,8 @@ class ChromiumDownloadSetter(DownloadSetter):
 
     def wait_download_begin(self, timeout: float = None) -> bool: ...
 
+    def wait_download_finish(self, timeout: float = None) -> bool: ...
+
     def _cookies_to_session(self) -> None: ...
 
     def _download_by_DownloadKit(self, **kwargs) -> None: ...

DrissionPage/web_page.py

@@ -555,7 +555,7 @@ class WebPageDownloadSetter(ChromiumDownloadSetter):
             try:
                 self._page.browser_driver.Browser.setDownloadBehavior(behavior='deny', eventsEnabled=True)
                 self._page.browser_driver.Browser.downloadWillBegin = self._download_by_DownloadKit
-                self._page.browser_driver.Browser.downloadProgress = None
+                # self._page.browser_driver.Browser.downloadProgress = None
             except CallMethodError:
                 raise RuntimeError('您的浏览器版本太低，不支持此方法，请升级。')