改进浏览器下载功能，可设定文件存在时处理方式

DrissionPage/chromium_page.py

@@ -15,7 +15,7 @@ from .chromium_base import ChromiumBase, Timeout, ChromiumBaseSetter, ChromiumBa
 from .chromium_driver import ChromiumDriver
 from .chromium_tab import ChromiumTab
 from .commons.browser import connect_browser
-from .commons.tools import port_is_using
+from .commons.tools import port_is_using, get_usable_path
 from .commons.web import set_session_cookies
 from .configs.chromium_options import ChromiumOptions
 from .errors import CallMethodError, BrowserConnectError
@@ -480,11 +480,13 @@ class ChromiumDownloadSetter(DownloadSetter):
         :param page: ChromiumPage对象
         """
         super().__init__(page)
-        self._behavior = 'allow'
+        self._behavior = 'allowAndName'
         self._download_th = None
         self._session = None
+        self._save_path = ''
         self._waiting_download = False
         self._download_begin = False
+        self._browser_missions = {}
 
     @property
     def session(self):
@@ -508,6 +510,7 @@ class ChromiumDownloadSetter(DownloadSetter):
         path = Path(path).absolute()
         path.mkdir(parents=True, exist_ok=True)
         path = str(path)
+        self._save_path = path
         self._page._download_path = path
         try:
             self._page.browser_driver.Browser.setDownloadBehavior(behavior='allowAndName', downloadPath=path,
@@ -523,20 +526,24 @@ class ChromiumDownloadSetter(DownloadSetter):
         try:
             self._page.browser_driver.Browser.setDownloadBehavior(behavior='allowAndName', eventsEnabled=True,
                                                                   downloadPath=self._page.download_path)
-            self._page.browser_driver.Browser.downloadWillBegin = self._download_by_browser
+            self._page.browser_driver.Browser.downloadWillBegin = self._download_will_begin
+            self._page.browser_driver.Browser.downloadProgress = self._download_progress
         except CallMethodError:
             self._page.driver.Page.setDownloadBehavior(behavior='allowAndName', downloadPath=self._page.download_path)
-            self._page.driver.Page.downloadWillBegin = self._download_by_browser
+            self._page.driver.Page.downloadWillBegin = self._download_will_begin
+            self._page.driver.Page.downloadProgress = self._download_progress
 
-        self._behavior = 'allow'
+        self._behavior = 'allowAndName'
 
     def by_DownloadKit(self):
         """设置使用DownloadKit下载文件"""
         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
         except CallMethodError:
             raise RuntimeError('您的浏览器版本太低，不支持此方法，请升级。')
+
         self._behavior = 'deny'
 
     def wait_download_begin(self, timeout=None):
@@ -583,16 +590,55 @@ class ChromiumDownloadSetter(DownloadSetter):
         if self._waiting_download:
             self._download_begin = True
 
-    def _download_by_browser(self, **kwargs):
+    def _download_will_begin(self, **kwargs):
-        """使用浏览器下载时调用"""
+        """浏览器下载即将开始时调用"""
+        m = BrowserDownloadMission(kwargs['guid'], kwargs['url'], kwargs['suggestedFilename'])
+        self._browser_missions[kwargs['guid']] = m
+        if self._file_exists == 'skip' and (Path(self._save_path) / kwargs["suggestedFilename"]).exists():
+            self._page.browser_driver.call_method('Browser.cancelDownload', guid=kwargs['guid'])
+            m.state = 'skipped'
+            p = Path(self._save_path) / kwargs["guid"]
+            if p.exists():
+                p.unlink()
+            return
+
         if self._waiting_download:
             self._download_begin = True
 
+    def _download_progress(self, **kwargs):
+        """下载状态产生变化时调用"""
+        guid = kwargs['guid']
+        m = self._browser_missions[guid]
+        m.size = kwargs['totalBytes']
+        m.received = kwargs['receivedBytes']
+        m.state = kwargs['state']
+
+        if m.state == 'completed':
+            path = Path(self._save_path) / m.name
+            from_path = Path(self._save_path) / guid
+            if path.exists():
+                if self._file_exists == 'rename':
+                    path = get_usable_path(path)
+                else:  # 'overwrite'
+                    path.unlink()
+            from_path.rename(path)
+
     def _wait_download_complete(self):
-        """等待下载完成"""
+        """等待DownloadKit下载完成"""
         self._page.download.wait()
 
 
+class BrowserDownloadMission(object):
+    def __init__(self, guid, url, name):
+        self.id = guid
+        self.url = url
+        self.name = name
+        self.save_path = None
+        self.state = None
+        self.size = None
+        self.received = None
+
+
 class Alert(object):
     """用于保存alert信息的类"""
 

DrissionPage/chromium_page.pyi

@@ -6,7 +6,7 @@
 from os import popen
 from pathlib import Path
 from threading import Thread
-from typing import Union, Tuple, List
+from typing import Union, Tuple, List, Dict
 
 from DownloadKit import DownloadKit
 from requests import Session
@@ -157,8 +157,10 @@ class ChromiumDownloadSetter(DownloadSetter):
         self._behavior: str = ...
         self._download_th: Thread = ...
         self._session: Session = None
+        self._save_path: str = None
         self._waiting_download: bool = ...
         self._download_begin: bool = ...
+        self._browser_missions: Dict[str, BrowserDownloadMission] = ...
 
     @property
     def session(self) -> Session: ...
@@ -178,11 +180,24 @@ class ChromiumDownloadSetter(DownloadSetter):
 
     def _download_by_DownloadKit(self, **kwargs) -> None: ...
 
-    def _download_by_browser(self, **kwargs) -> None: ...
+    def _download_will_begin(self, **kwargs) -> None: ...
+
+    def _download_progress(self, **kwargs) -> None: ...
 
     def _wait_download_complete(self) -> None: ...
 
 
+class BrowserDownloadMission(object):
+    def __init__(self, guid: str, url: str, name: str):
+        self.id: str = ...
+        self.url: str = ...
+        self.name: str = ...
+        self.save_path: str = ...
+        self.state: str = ...
+        self.size: str = ...
+        self.received: str = ...
+
+
 class Alert(object):
 
     def __init__(self):

DrissionPage/web_page.py

@@ -518,6 +518,7 @@ class WebPageDownloadSetter(ChromiumDownloadSetter):
         path = Path(path).absolute()
         path.mkdir(parents=True, exist_ok=True)
         path = str(path)
+        self._save_path = path
         self._page._download_path = path
         self.DownloadKit.goal_path = path
 
@@ -537,14 +538,16 @@ class WebPageDownloadSetter(ChromiumDownloadSetter):
         try:
             self._page.browser_driver.Browser.setDownloadBehavior(behavior='allowAndName', eventsEnabled=True,
                                                                   downloadPath=self._page.download_path)
-            self._page.browser_driver.Browser.downloadWillBegin = self._download_by_browser
+            self._page.browser_driver.Browser.downloadWillBegin = self._download_will_begin
+            self._page.browser_driver.Browser.downloadProgress = self._download_progress
 
         except CallMethodError:
             warn('\n您的浏览器版本太低，用新标签页下载文件可能崩溃，建议升级。')
             self._page.driver.Page.setDownloadBehavior(behavior='allowAndName', downloadPath=self._page.download_path)
-            self._page.driver.Page.downloadWillBegin = self._download_by_browser
+            self._page.driver.Page.downloadWillBegin = self._download_will_begin
+            self._page.driver.Page.downloadProgress = self._download_progress
 
-        self._behavior = 'allow'
+        self._behavior = 'allowAndName'
 
     def by_DownloadKit(self):
         """设置使用DownloadKit下载文件"""
@@ -552,6 +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
             except CallMethodError:
                 raise RuntimeError('您的浏览器版本太低，不支持此方法，请升级。')