diff --git a/src/managers/driver/WebDriverDownloader.py b/src/managers/driver/WebDriverDownloader.py index de3bf74..a3d9003 100644 --- a/src/managers/driver/WebDriverDownloader.py +++ b/src/managers/driver/WebDriverDownloader.py @@ -1,6 +1,7 @@ import os import time import shutil +import threading import requests import zipfile import tarfile @@ -243,29 +244,33 @@ class WebDriverDownloader: def download( self, - progress_callback: Optional[Callable[[int, int, float, str], None]] = None + progress_callback: Optional[Callable[[float, int, float, str], None]] = None, + cancel_event: Optional[threading.Event] = None ) -> Optional[Path]: try: # downlaod file : 0% - 98% - if not self._download(progress_callback): + if not self._download(progress_callback, cancel_event=cancel_event): return None # verify file : 98% - 99% if not self._verify(progress_callback): + progress_callback(0, 100, 0.0, "验证失败") return None # extract file : 99% - 100% driver_path = self._extract(progress_callback) if not driver_path: + progress_callback(0, 100, 0.0, "解压失败") return None return driver_path - except Exception: - return None + except Exception as e: + raise e def _download( self, - progress_callback: Optional[Callable[[int, int, float, str], None]] = None, - max_retries: int = 3 + progress_callback: Optional[Callable[[float, int, float, str], None]] = None, + max_retries: int = 3, + cancel_event: Optional[threading.Event] = None ) -> bool: CHUNK_SIZE = 8192*8 # 64KB chunk @@ -276,6 +281,8 @@ class WebDriverDownloader: for attempt in range(max_retries): try: + if cancel_event and cancel_event.is_set(): + return False # resume download if file exists if self.download_path.exists(): downloaded_size = self.download_path.stat().st_size @@ -287,7 +294,7 @@ class WebDriverDownloader: headers_ = headers mode = 'wb' # get response - response = requests.get(str(self.download_url), headers=headers_, stream=True, timeout=120) + response = requests.get(str(self.download_url), headers=headers_, stream=True, timeout=10) if response.status_code not in [200, 206]: if self.download_path.exists(): self.download_path.unlink() @@ -306,6 +313,9 @@ class WebDriverDownloader: last_progress = 0.0 with open(self.download_path, mode) as f: for chunk in response.iter_content(CHUNK_SIZE): + if cancel_event and cancel_event.is_set(): + response.close() + return False if not chunk: continue f.write(chunk) @@ -328,8 +338,10 @@ class WebDriverDownloader: raise Exception(f"下载不完整 : {self.download_path.stat().st_size}/{total_size} 字节") return True except Exception as e: + if cancel_event and cancel_event.is_set(): + return False if attempt < max_retries - 1: - progress_callback(0, 100, 0.0, "准备重试...") + progress_callback(0, 100, 0.0, f"第 {attempt+1} 次重试...") time.sleep(1) continue raise e @@ -337,7 +349,7 @@ class WebDriverDownloader: def _verify( self, - progress_callback: Optional[Callable[[int, int, float, str], None]] = None + progress_callback: Optional[Callable[[float, int, float, str], None]] = None ) -> bool: progress_callback(98, 100, 0.0, "验证完成") @@ -346,7 +358,7 @@ class WebDriverDownloader: def _extract( self, - progress_callback: Optional[Callable[[int, int, float, str], None]] = None + progress_callback: Optional[Callable[[float, int, float, str], None]] = None ) -> Optional[Path]: try: diff --git a/src/managers/driver/WebDriverManager.py b/src/managers/driver/WebDriverManager.py index afd44fe..3e70bd8 100644 --- a/src/managers/driver/WebDriverManager.py +++ b/src/managers/driver/WebDriverManager.py @@ -40,21 +40,22 @@ class WebDriverInfo: Web driver information. Attributes: - browser_info (WebBrowserInfo): Web browser information driver_type (WebDriverType): Web driver type + driver_arch (WebDriverArch): Web driver architecture driver_version (str): Web driver version + browser_version (str): Web browser version driver_path (Optional[Path]): Web driver executable file path driver_status (DriverStatus): Web driver status """ def __init__( - self, - browser_info: WebBrowserInfo + self ): - self.browser_info = browser_info - self.driver_type = WebDriverType(browser_info.browser_type.value) + self.driver_type = None + self.driver_arch = None self.driver_version = "" + self.browser_version = "" self.driver_path: Optional[Path] = None self.driver_status = DriverStatus.NOT_INSTALLED @@ -99,7 +100,10 @@ class WebDriverManager: with self.__lock: browser_infos = self.__browser_detector.detect() - self.__driver_infos = [WebDriverInfo(info) for info in browser_infos] + self.__driver_infos = [ + self._getDriverInfo(info) + for info in browser_infos + ] def _checkDriverStatus( @@ -108,27 +112,28 @@ class WebDriverManager: with self.__lock: for driver_info in self.__driver_infos: - driver_arch = self._mapWebBrowserArch( - driver_info.browser_info.browser_type, - driver_info.browser_info.browser_arch - ) - driver_path = self._getDriverPath( - driver_info.driver_type, - driver_arch - ) + driver_path = self._getDriverPath(driver_info) if driver_path and driver_path.exists() and driver_path.is_file(): driver_info.driver_path = driver_path driver_info.driver_status = DriverStatus.INSTALLED - try: - driver_info.driver_version = self._getDriverVersion( - driver_info.driver_type, - driver_info.driver_info.browser_version - ) - except Exception: - driver_info.driver_status = DriverStatus.ERROR - def _mapWebBrowserArch( + def _mapWebBrowserTypeToDriver( + self, + browser_type: WebBrowserType + ) -> WebDriverType: + + if browser_type == WebBrowserType.CHROME: + return WebDriverType.CHROME + elif browser_type == WebBrowserType.FIREFOX: + return WebDriverType.FIREFOX + elif browser_type == WebBrowserType.EDGE: + return WebDriverType.EDGE + else: + raise ValueError(f"不支持的 Web 浏览器类型 : {browser_type}") + + + def _mapWebBrowserArchToDriver( self, browser_type: WebBrowserType, browser_arch: WebBrowserArch @@ -236,12 +241,30 @@ class WebDriverManager: raise ValueError(f"无效的 Firefox 版本格式 : {version}") from e + def _getDriverInfo( + self, + browser_info: WebBrowserInfo + ) -> WebDriverInfo: + + driver_info = WebDriverInfo() + driver_info.driver_type = self._mapWebBrowserTypeToDriver(browser_info.browser_type) + driver_info.driver_arch = self._mapWebBrowserArchToDriver(browser_info.browser_type, browser_info.browser_arch) + if browser_info.browser_type == WebBrowserType.FIREFOX: + driver_info.driver_version = self._mapFirefoxDriverVersion(browser_info.browser_version) + else: + driver_info.driver_version = browser_info.browser_version + driver_info.browser_version = browser_info.browser_version + return driver_info + + def _getDriverPath( self, - driver_type: WebDriverType, - driver_arch: WebDriverArch + driver_info: WebDriverInfo ) -> Optional[Path]: + driver_type = driver_info.driver_type + driver_arch = driver_info.driver_arch + driver_version = driver_info.driver_version if driver_type == WebDriverType.CHROME: driver_name = "chromedriver" elif driver_type == WebDriverType.FIREFOX: @@ -259,31 +282,17 @@ class WebDriverManager: WebDriverArch.Edge.WINX86_64, ] exe_name = f"{driver_name}.exe" if is_win else driver_name - driver_dir = Path(self.__driver_dir) / driver_type.value / driver_arch.value - driver_path = driver_dir / exe_name - if driver_path.exists() and driver_path.is_file(): - return driver_path - return None - - - def _getDriverVersion( - self, - driver_type: WebDriverType, - browser_version: str - ) -> str: - - if driver_type == WebDriverType.FIREFOX: - return self._mapFirefoxDriverVersion(browser_version) - return browser_version + driver_dir = Path(self.__driver_dir)/driver_type.value/driver_version/driver_arch.value + driver_path = driver_dir/exe_name + return driver_path def refresh( self ): - with self.__lock: - self._detectBrowsers() - self._checkDriverStatus() + self._detectBrowsers() + self._checkDriverStatus() def getDriverInfos( @@ -297,21 +306,21 @@ class WebDriverManager: def getDriverInfo( self, driver_type: WebDriverType - ) -> Optional[WebDriverInfo]: + ) -> list[WebDriverInfo]: with self.__lock: - for driver_info in self.__driver_infos: - if driver_info.driver_type == driver_type: - return driver_info - return None + return [ + info + for info in self.__driver_infos + if info.driver_type == driver_type + ] def getDriverPath( self, - driver_type: WebDriverType + driver_info: WebDriverInfo ) -> Optional[Path]: - driver_info = self.getDriverInfo(driver_type) if driver_info and driver_info.driver_status == DriverStatus.INSTALLED: return driver_info.driver_path return None @@ -319,24 +328,28 @@ class WebDriverManager: def installDriver( self, - driver_type: WebDriverType, - progress_callback: Optional[Callable[[int, int, float, str], None]] = None + driver_info: WebDriverInfo, + progress_callback: Optional[Callable[[float, int, float, str], None]] = None, + cancel_event: Optional[threading.Event] = None ) -> Optional[Path]: with self.__lock: - driver_info = self.getDriverInfo(driver_type) if not driver_info: - raise ValueError(f"未找到类型为 {driver_type} 的浏览器") - if driver_info.driver_status == DriverStatus.DOWNLOADING: - raise ValueError(f"{driver_type} 驱动正在下载中") - driver_info.driver_status = DriverStatus.DOWNLOADING + if progress_callback: + progress_callback(0, 0, 0, "未找到浏览器信息") + else: + raise ValueError("未找到浏览器信息") + if driver_info and driver_info.driver_status == DriverStatus.DOWNLOADING: + if progress_callback: + progress_callback(0, 0, 0, f"{driver_info.driver_type} 驱动正在下载中") + else: + raise ValueError(f"{driver_info.driver_type} 驱动正在下载中") try: - driver_arch = self._mapWebBrowserArch( - driver_info.browser_info.browser_type, - driver_info.browser_info.browser_arch - ) - browser_version = driver_info.browser_info.browser_version - driver_version = self._getDriverVersion(driver_type, browser_version) + if not driver_info: + raise ValueError("未找到浏览器信息") + driver_arch = driver_info.driver_arch + driver_type = driver_info.driver_type + driver_version = driver_info.driver_version downloader = None if driver_type == WebDriverType.CHROME: downloader = ChromeDriverDownloader( @@ -357,9 +370,13 @@ class WebDriverManager: download_dir=self.__driver_dir ) if downloader is None: - raise ValueError(f"不支持的 Web Driver 类型 : {driver_type}") - - driver_path = downloader.download(progress_callback=progress_callback) + if progress_callback: + progress_callback(0, 0, 0, f"不支持的 Web Driver 类型") + else: + raise ValueError(f"不支持的 Web Driver 类型") + with self.__lock: + driver_info.driver_status = DriverStatus.DOWNLOADING + driver_path = downloader.download(progress_callback=progress_callback, cancel_event=cancel_event) with self.__lock: if driver_path: driver_info.driver_path = driver_path @@ -369,6 +386,57 @@ class WebDriverManager: driver_info.driver_status = DriverStatus.ERROR return driver_path except Exception as e: + with self.__lock: + driver_info.driver_status = DriverStatus.ERROR + raise e + + + def cancelDriverDownload( + self, + driver_info: WebDriverInfo + ) -> bool: + + import shutil + + try: + driver_path = self._getDriverPath(driver_info) + if driver_path: + download_dir = driver_path.parent + if download_dir.exists(): + shutil.rmtree(download_dir, ignore_errors=True) + with self.__lock: + driver_info.driver_path = None + driver_info.driver_status = DriverStatus.NOT_INSTALLED + return True + except Exception: + return False + + + def uninstallDriver( + self, + driver_info: WebDriverInfo, + progress_callback: Optional[Callable[[int, int, float, str], None]] = None + ) -> bool: + + with self.__lock: + if not driver_info: + if progress_callback: + progress_callback(0, 0, 0, "未找到浏览器信息") + else: + raise ValueError("未找到浏览器信息") + if driver_info.driver_status != DriverStatus.INSTALLED: + if progress_callback: + progress_callback(0, 0, 0, f"{driver_info.driver_type} 驱动未安装") + else: + raise ValueError(f"{driver_info.driver_type} 驱动未安装") + try: + driver_path = driver_info.driver_path + driver_path.unlink() + with self.__lock: + driver_info.driver_path = None + driver_info.driver_status = DriverStatus.NOT_INSTALLED + return True + except Exception: with self.__lock: driver_info.driver_status = DriverStatus.ERROR raise @@ -399,5 +467,5 @@ def instance( _webdriver_manager_instance = WebDriverManager(driver_dir) else: if driver_dir and _webdriver_manager_instance.driverDir() != os.path.abspath(driver_dir): - raise ValueError("WebDriverManager 的实例已初始化,不能使用不同的驱动目录") + raise ValueError("WebDriverManager 的实例已初始化, 不能使用不同的驱动目录") return _webdriver_manager_instance