Pre Merge pull request !11 from 熊猫同学/develop

DrissionPage/easy_set.py

@@ -5,12 +5,12 @@
 @File    :   driver_page.py
 """
 from os import popen
+
 from pathlib import Path
 from pprint import pprint
 from re import search as RE_SEARCH
-from typing import Union
-
 from selenium import webdriver
+from typing import Union
 
 from DrissionPage.config import OptionsManager, DriverOptions
 from DrissionPage.drission import Drission
@@ -190,27 +190,36 @@ def check_driver_version(driver_path: str = None, chrome_path: str = None) -> bo
 
 
 # -------------------------自动识别chrome版本号并下载对应driver------------------------
-def get_match_driver(ini_path: str = None,
+def get_match_driver(ini_path: Union[str, None] = 'default',
                      save_path: str = None,
-                     chrome_path: str = None) -> None:
+                     chrome_path: str = None,
+                     show_msg: bool = True,
+                     check_version: bool = True) -> Union[str, None]:
     """自动识别chrome版本并下载匹配的driver             \n
     :param ini_path: 要读取和修改的ini文件路径
     :param save_path: chromedriver保存路径
     :param chrome_path: 指定chrome.exe位置
+    :param show_msg: 是否打印信息
+    :param check_version: 是否检查版本匹配
     :return: None
     """
-    save_path = save_path or str(Path(__file__).parent)
+    # save_path = save_path or str(Path(__file__).parent)
 
-    chrome_path = chrome_path or _get_chrome_path(ini_path)
+    chrome_path = chrome_path or _get_chrome_path(ini_path, show_msg)
-    chrome_path = Path(chrome_path).absolute() if chrome_path else None
+    chrome_absolute_path = Path(chrome_path).absolute()
-    print('chrome.exe路径', chrome_path, '\n')
+    # 默认将文件直接放在浏览器同目录，方便查找
+    save_path = save_path or str(chrome_absolute_path.parent)
+    chrome_path = chrome_absolute_path if chrome_path else None
+    if show_msg:
+        print('chrome.exe路径', chrome_path)
 
-    ver = _get_chrome_version(chrome_path)
+    ver = _get_chrome_version(str(chrome_path))
-    print('version', ver, '\n')
+    if show_msg:
+        print('version', ver)
 
-    zip_path = _download_driver(ver, save_path)
+    zip_path = _download_driver(ver, save_path, show_msg=show_msg)
 
-    if not zip_path:
+    if not zip_path and show_msg:
         print('没有找到对应版本的driver。')
 
     try:
@@ -218,52 +227,89 @@ def get_match_driver(ini_path: str = None,
     except TypeError:
         driver_path = None
 
-    print('\n解压路径', driver_path, '\n')
+    if show_msg:
+        print('解压路径', driver_path)
 
     if driver_path:
         Path(zip_path).unlink()
-        set_paths(driver_path=driver_path, chrome_path=str(chrome_path), ini_path=ini_path, check_version=False)
+        if ini_path:
+            set_paths(driver_path=driver_path, chrome_path=str(chrome_path), ini_path=ini_path, check_version=False)
 
-        if not check_driver_version(driver_path, chrome_path):
+        if check_version:
-            print('获取失败，请手动配置。')
+            if not check_driver_version(driver_path, chrome_path) and show_msg:
+                print('获取失败，请手动配置。')
     else:
-        print('获取失败，请手动配置。')
+        if show_msg:
+            print('获取失败，请手动配置。')
+
+    return driver_path
 
 
-def _get_chrome_path(ini_path: str = None) -> Union[str, None]:
+def _get_chrome_path(ini_path: str = None,
+                     show_msg: bool = True,
+                     from_ini: bool = True,
+                     from_regedit: bool = True,
+                     from_system_path: bool = True, ) -> Union[str, None]:
     """从ini文件或系统变量中获取chrome.exe的路径    \n
     :param ini_path: ini文件路径
     :return: chrome.exe路径
     """
     # -----------从ini文件中获取--------------
-    try:
+    if ini_path and from_ini:
-        path = OptionsManager(ini_path).chrome_options['binary_location']
+        try:
-    except KeyError:
+            path = OptionsManager(ini_path).chrome_options['binary_location']
-        return None
+        except KeyError:
+            path = None
+    else:
+        path = None
 
     if path and Path(path).is_file():
         print('ini文件中', end='')
         return str(path)
 
-    # -----------从系统路径中获取--------------
+    # -----------从注册表中获取--------------
-    paths = popen('set path').read().lower()
+    if from_regedit:
-    r = RE_SEARCH(r'[^;]*chrome[^;]*', paths)
+        import winreg
+        try:
+            key = winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE,
+                                 r'SOFTWARE\Microsoft\Windows\CurrentVersion\App Paths\chrome.exe',
+                                 reserved=0, access=winreg.KEY_READ)
+            k = winreg.EnumValue(key, 0)
+            winreg.CloseKey(key)
 
-    if r:
+            if show_msg:
-        path = Path(r.group(0)) if 'chrome.exe' in r.group(0) else Path(r.group(0)) / 'chrome.exe'
+                print('注册表中', end='')
 
-        if path.exists():
+            return k[1]
-            print('系统中', end='')
-            return str(path)
 
-    paths = paths.split(';')
+        except FileNotFoundError:
+            pass
 
-    for path in paths:
+    # -----------从系统变量中获取--------------
-        path = Path(path) / 'chrome.exe'
+    if from_system_path:
+        paths = popen('set path').read().lower()
+        r = RE_SEARCH(r'[^;]*chrome[^;]*', paths)
 
-        if path.exists():
+        if r:
-            print('系统中', end='')
+            path = Path(r.group(0)) if 'chrome.exe' in r.group(0) else Path(r.group(0)) / 'chrome.exe'
-            return str(path)
+
+            if path.exists():
+                if show_msg:
+                    print('系统变量中', end='')
+                return str(path)
+
+        paths = paths.split(';')
+
+        for path in paths:
+            path = Path(path) / 'chrome.exe'
+
+            try:
+                if path.exists():
+                    if show_msg:
+                        print('系统变量中', end='')
+                    return str(path)
+            except OSError:
+                pass
 
 
 def _get_chrome_version(path: str) -> Union[str, None]:
@@ -283,7 +329,7 @@ def _get_chrome_version(path: str) -> Union[str, None]:
         return None
 
 
-def _download_driver(version: str, save_path: str = None) -> Union[str, None]:
+def _download_driver(version: str, save_path: str = None, show_msg: bool = True) -> Union[str, None]:
     """根据传入的版本号到镜像网站查找，下载最相近的          \n
     :param version: 本地版本号
     :return: 保存地址
@@ -298,26 +344,42 @@ def _download_driver(version: str, save_path: str = None) -> Union[str, None]:
     loc_main = version.split('.')[0]
 
     try:
-        loc_num = int(version.replace('.', ''))
+        # 改变查询逻辑，对应大版本号即可
+        loc_num = ''.join(version.split('.')[:3])
     except ValueError:
         return None
+    # 将旧版本和新版本链接分开，提高遍历效率
+    old_versions = page.eles('xpath://pre/a')[:48]
+    new_versions = page.eles('xpath://pre/a')[48:]
+    versions = []
+    if loc_num == '2':
+        versions = old_versions
+    else:
+        versions = new_versions
 
-    for i in page.eles('xpath://pre/a'):
+    for i in versions:
         remote_main = i.text.split('.')[0]
-
+        # 大版本号不同则跳过循环
-        try:
+        if remote_main != loc_main:
-            remote_num = int(i.text.replace('.', '').replace('/', ''))
+            # 碰到非版本目录停止遍历
-        except ValueError:
+            if remote_main == 'icons/':
+                if show_msg:
+                    print('找到最相近的版本为：{}'.format(remote_ver.replace('/', '')))
+                break
             continue
 
-        if remote_main == loc_main and remote_num >= loc_num:
+        try:
+            remote_num = i.text.replace('.', '').replace('/', '')
+        except ValueError:
+            continue
+        # 碰到相同版本号取最新的
+        if loc_num in remote_num:
             remote_ver = i.text
-            break
 
     if remote_ver:
         url = f'https://cdn.npm.taobao.org/dist/chromedriver/{remote_ver}chromedriver_win32.zip'
         save_path = save_path or Path(__file__).parent
-        result = page.download(url, save_path, file_exists='overwrite', show_msg=True)
+        result = page.download(url, save_path, file_exists='overwrite', show_msg=show_msg)
 
         if result[0]:
             return result[1]