From ed2883e2b96ca75f1d0774c1263a813a3c38cded Mon Sep 17 00:00:00 2001
From: g1879 <g1879@qq.com>
Date: Mon, 11 Sep 2023 07:10:26 +0800
Subject: [PATCH] =?UTF-8?q?browser=5Fdriver=E5=92=8CBrowserManager?=
 =?UTF-8?q?=E6=94=B9=E6=88=90=E6=AF=8F=E4=B8=AA=E6=B5=8F=E8=A7=88=E5=99=A8?=
 =?UTF-8?q?=E5=8F=AA=E4=B8=80=E4=B8=AA=E5=AF=B9=E8=B1=A1=EF=BC=9B=E6=94=B9?=
 =?UTF-8?q?=E8=BF=9B=E4=B8=8B=E6=B5=8F=E8=A7=88=E5=99=A8=E8=BD=BD=E5=8A=9F?=
 =?UTF-8?q?=E8=83=BD=EF=BC=8C=E6=9C=AA=E5=AE=8C=E6=88=90?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 DrissionPage/chromium_driver.py  | 18 +++++++++
 DrissionPage/chromium_driver.pyi |  6 ++-
 DrissionPage/chromium_page.py    | 69 +++++++++++++++++++++++---------
 DrissionPage/chromium_page.pyi   |  7 +++-
 DrissionPage/waiter.py           | 17 +-------
 DrissionPage/waiter.pyi          |  2 -
 6 files changed, 80 insertions(+), 39 deletions(-)

diff --git a/DrissionPage/chromium_driver.py b/DrissionPage/chromium_driver.py
index fde230f..b638a0e 100644
--- a/DrissionPage/chromium_driver.py
+++ b/DrissionPage/chromium_driver.py
@@ -233,3 +233,21 @@ class ChromiumDriver(object):
         return f"<ChromiumDriver {self.id}>"
 
     __repr__ = __str__
+
+
+class BrowserDriver(ChromiumDriver):
+    BROWSERS = {}
+
+    def __new__(cls, tab_id, tab_type, address):
+        if tab_id in cls.BROWSERS:
+            return cls.BROWSERS[tab_id]
+        return object.__new__(cls)
+
+    def __init__(self, tab_id, tab_type, address):
+        if tab_id in BrowserDriver.BROWSERS:
+            return
+        super().__init__(tab_id, tab_type, address)
+        BrowserDriver.BROWSERS[tab_id] = self
+
+    def __repr__(self):
+        return f"<BrowserDriver {self.id}>"
diff --git a/DrissionPage/chromium_driver.pyi b/DrissionPage/chromium_driver.pyi
index 9ba7feb..152d4c4 100644
--- a/DrissionPage/chromium_driver.pyi
+++ b/DrissionPage/chromium_driver.pyi
@@ -5,7 +5,7 @@
 """
 from queue import Queue
 from threading import Thread, Event
-from typing import Union, Callable
+from typing import Union, Callable, Dict
 
 
 class GenericAttr(object):
@@ -58,3 +58,7 @@ class ChromiumDriver(object):
     def get_listener(self, event: str) -> Union[Callable, None]: ...
 
     def __str__(self) -> str: ...
+
+
+class BrowserDriver(ChromiumDriver):
+    BROWSERS: Dict[str, ChromiumDriver] = ...
diff --git a/DrissionPage/chromium_page.py b/DrissionPage/chromium_page.py
index 8009d16..742f05c 100644
--- a/DrissionPage/chromium_page.py
+++ b/DrissionPage/chromium_page.py
@@ -3,12 +3,14 @@
 @Author  :   g1879
 @Contact :   g1879@qq.com
 """
+from pathlib import Path
 from shutil import move
+from threading import Lock
 from time import perf_counter, sleep
 
 from .chromium_base import ChromiumBase, Timeout
 from .chromium_base import handle_download
-from .chromium_driver import ChromiumDriver
+from .chromium_driver import ChromiumDriver, BrowserDriver
 from .chromium_tab import ChromiumTab
 from .commons.browser import connect_browser
 from .commons.tools import get_usable_path
@@ -98,7 +100,7 @@ class ChromiumPage(ChromiumBase):
         u = f'http://{self.address}/json/version'
         ws = self._control_session.get(u).json()['webSocketDebuggerUrl']
         self._control_session.get(u, headers={'Connection': 'close'})
-        self._browser_driver = ChromiumDriver(ws.split('/')[-1], 'browser', self.address)
+        self._browser_driver = BrowserDriver(ws.split('/')[-1], 'browser', self.address)
         self._browser_driver.start()
 
         self._alert = Alert()
@@ -446,13 +448,32 @@ class ChromiumTabRect(object):
 
 
 class BrowserDownloadManager(object):
+    BROWSERS = {}
+
+    def __new__(cls, page):
+        """
+        :param page: ChromiumPage对象
+        """
+        if page.browser_driver.id in cls.BROWSERS:
+            return cls.BROWSERS[page.browser_driver.id]
+        return object.__new__(cls)
+
     def __init__(self, page):
+        """
+        :param page: ChromiumPage对象
+        """
+        if page.browser_driver.id in BrowserDownloadManager.BROWSERS:
+            return
+
         self._page = page
+        self._lock = Lock()
         page.set.download_path(page.download_path)
         self._page.browser_driver.set_listener('Browser.downloadProgress', self._onDownloadProgress)
         self._page.browser_driver.set_listener('Browser.downloadWillBegin', self._onDownloadWillBegin)
         self._missions = {}
 
+        BrowserDownloadManager.BROWSERS[page.browser_driver.id] = self
+
     @property
     def missions(self):
         return self._missions
@@ -464,12 +485,20 @@ class BrowserDownloadManager(object):
         """
         self._missions[mission.id] = mission
 
-    def cancel(self, mission):
-        """取消一个下载任务
+    def set_done(self, mission, state, cancel=False, final_path=None):
+        """设置任务结束
         :param mission: 任务对象
+        :param state: 任务状态
+        :param cancel: 是否取消
+        :param final_path: 最终路径
         :return: None
         """
-        self._page.browser_driver.call_method('Browser.cancelDownload', guid=mission.id)
+        mission.state = state
+        mission.final_path = final_path
+        if cancel:
+            self._page.browser_driver.call_method('Browser.cancelDownload', guid=mission.id)
+            if mission.final_path:
+                Path(mission.final_path).unlink(True)
         self._missions.pop(mission.id)
 
     def _onDownloadWillBegin(self, **kwargs):
@@ -481,22 +510,24 @@ class BrowserDownloadManager(object):
     def _onDownloadProgress(self, **kwargs):
         """下载状态变化时执行"""
         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']
+            with self._lock:
+                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':
-                mission.received_bytes = kwargs['receivedBytes']
-                mission.total_bytes = kwargs['totalBytes']
-                form_path = f'{self._page.download_path}\\{mission.id}'
-                to_path = str(get_usable_path(f'{mission.path}\\{mission.name}'))
-                move(form_path, to_path)
-                mission._set_done('completed', final_path=to_path)
+                    elif kwargs['state'] == 'completed':
+                        mission.received_bytes = kwargs['receivedBytes']
+                        mission.total_bytes = kwargs['totalBytes']
+                        form_path = f'{self._page.download_path}\\{mission.id}'
+                        to_path = str(get_usable_path(f'{mission.path}\\{mission.name}'))
+                        move(form_path, to_path)
+                        self.set_done(mission, 'completed', final_path=to_path)
 
-            else:
-                mission._set_done('canceled')
+                    else:
+                        self.set_done(mission, 'canceled')
 
 
 class Alert(object):
diff --git a/DrissionPage/chromium_page.pyi b/DrissionPage/chromium_page.pyi
index be96790..88d62e0 100644
--- a/DrissionPage/chromium_page.pyi
+++ b/DrissionPage/chromium_page.pyi
@@ -3,6 +3,7 @@
 @Author  :   g1879
 @Contact :   g1879@qq.com
 """
+from threading import Lock
 from typing import Union, Tuple, List, Dict
 
 from .chromium_base import ChromiumBase
@@ -128,6 +129,10 @@ class ChromiumTabRect(object):
 class BrowserDownloadManager(object):
     _page: ChromiumPage = ...
     _missions: Dict[str, DownloadMission] = ...
+    _lock: Lock = ...
+    BROWSERS: Dict[str, BrowserDownloadManager] = ...
+
+    def __new__(cls, page: ChromiumPage): ...
 
     def __init__(self, page: ChromiumPage): ...
 
@@ -136,7 +141,7 @@ class BrowserDownloadManager(object):
 
     def add_mission(self, mission: DownloadMission) -> None: ...
 
-    def cancel(self, mission: DownloadMission) -> None: ...
+    def set_done(self, mission: DownloadMission, state: str, cancel: bool = False, final_path: str = None) -> None: ...
 
     def _onDownloadWillBegin(self, **kwargs) -> None: ...
 
diff --git a/DrissionPage/waiter.py b/DrissionPage/waiter.py
index c8f73a4..e693873 100644
--- a/DrissionPage/waiter.py
+++ b/DrissionPage/waiter.py
@@ -401,9 +401,7 @@ class DownloadMission(object):
 
     def cancel(self):
         """取消该任务，如任务已完成，删除已下载的文件"""
-        self._set_done('canceled', True)
-        if self.final_path:
-            Path(self.final_path).unlink(True)
+        self.tab._page._dl_mgr.set_done(self, state='canceled', cancel=True)
 
     def wait(self, show=True, timeout=None, cancel_if_timeout=True):
         """等待任务结束
@@ -450,16 +448,3 @@ class DownloadMission(object):
             print()
 
         return self.final_path if self.final_path else False
-
-    def _set_done(self, state, cancel=False, final_path=None):
-        """设置任务结束
-        :param state: 任务状态
-        :param cancel: 是否取消
-        :param final_path: 最终路径
-        :return: None
-        """
-        self.state = state
-        self.final_path = final_path
-        if cancel:
-            self.tab._page._dl_mgr.cancel(self)
-        self.tab._download_missions.remove(self)
diff --git a/DrissionPage/waiter.pyi b/DrissionPage/waiter.pyi
index 7858606..ad73281 100644
--- a/DrissionPage/waiter.pyi
+++ b/DrissionPage/waiter.pyi
@@ -108,5 +108,3 @@ class DownloadMission(object):
     def cancel(self) -> None: ...
 
     def wait(self, show: bool = True, timeout=None, cancel_if_timeout=True) -> Union[bool, str]: ...
-
-    def _set_done(self, state: str, cancel: bool = False, final_path: str = None) -> None: ...