替换判断文件名有效的方法为get_usable_path()

DrissionPage/common.py

@@ -6,7 +6,7 @@
 """
 from html import unescape
 from pathlib import Path
-from re import split as re_SPLIT
+from re import split as re_SPLIT, search, sub
 from shutil import rmtree
 from typing import Union
 from zipfile import ZipFile
@@ -213,31 +213,6 @@ def translate_loc(loc: tuple) -> tuple:
     return loc_by, loc_str
 
 
-def get_available_file_name(folder_path: str, file_name: str) -> str:
-    """检查文件是否重名，并返回可以使用的文件名  \n
-    :param folder_path: 文件夹路径
-    :param file_name: 要检查的文件名
-    :return: 可用的文件名
-    """
-    folder_path = Path(folder_path).absolute()
-    file_Path = folder_path.joinpath(file_name)
-
-    while file_Path.exists():
-        ext_name = file_Path.suffix
-        base_name = file_Path.stem
-        num = base_name.split(' ')[-1]
-
-        if num and num[0] == '(' and num[-1] == ')' and num[1:-1].isdigit():
-            num = int(num[1:-1])
-            file_name = f'{base_name.replace(f"({num})", "", -1)}({num + 1}){ext_name}'
-        else:
-            file_name = f'{base_name} (1){ext_name}'
-
-        file_Path = folder_path.joinpath(file_name)
-
-    return file_name
-
-
 def clean_folder(folder_path: str, ignore: list = None) -> None:
     """清空一个文件夹，除了ignore里的文件和文件夹  \n
     :param folder_path: 要清空的文件夹路径
@@ -284,3 +259,66 @@ def get_exe_path_from_port(port: Union[str, int]) -> Union[str, None]:
     else:
         file_lst = popen(f'wmic process where processid={processid} get executablepath').read().split('\n')
         return file_lst[2].strip() if len(file_lst) > 2 else None
+
+
+def get_usable_path(path: Union[str, Path]) -> Path:
+    """检查文件或文件夹是否有重名，并返回可以使用的路径           \n
+    :param path: 文件或文件夹路径
+    :return: 可用的路径，Path对象
+    """
+    path = Path(path)
+    parent = path.parent
+    path = parent / make_valid_name(path.name)
+    name = path.stem if path.is_file() else path.name
+    ext = path.suffix if path.is_file() else ''
+
+    first_time = True
+
+    while path.exists():
+        r = search(r'(.*)_(\d+)$', name)
+
+        if not r or (r and first_time):
+            src_name, num = name, '1'
+        else:
+            src_name, num = r.group(1), int(r.group(2)) + 1
+
+        name = f'{src_name}_{num}'
+        path = parent / f'{name}{ext}'
+        first_time = None
+
+    return path
+
+
+def make_valid_name(full_name: str) -> str:
+    """获取有效的文件名                  \n
+    :param full_name: 文件名
+    :return: 可用的文件名
+    """
+    # ----------------去除前后空格----------------
+    full_name = full_name.strip()
+
+    # ----------------使总长度不大于255个字符（一个汉字是2个字符）----------------
+    r = search(r'(.*)(\.[^.]+$)', full_name)  # 拆分文件名和后缀名
+    if r:
+        name, ext = r.group(1), r.group(2)
+        ext_long = len(ext)
+    else:
+        name, ext = full_name, ''
+        ext_long = 0
+
+    while get_long(name) > 255 - ext_long:
+        name = name[:-1]
+
+    full_name = f'{name}{ext}'
+
+    # ----------------去除不允许存在的字符----------------
+    return sub(r'[<>/\\|:*?\n]', '', full_name)
+
+
+def get_long(txt) -> int:
+    """返回字符串中字符个数（一个汉字是2个字符）          \n
+    :param txt: 字符串
+    :return: 字符个数
+    """
+    txt_len = len(txt)
+    return int((len(txt.encode('utf-8')) - txt_len) / 2 + txt_len)

DrissionPage/driver_element.py

@@ -16,7 +16,7 @@ from selenium.webdriver.support import expected_conditions as ec
 from selenium.webdriver.support.wait import WebDriverWait
 
 from .base import DrissionElement, BaseElement
-from .common import str_to_loc, get_available_file_name, translate_loc, format_html
+from .common import str_to_loc, get_usable_path, translate_loc, format_html
 from .session_element import make_session_ele
 
 
@@ -413,8 +413,8 @@ class DriverElement(DrissionElement):
         name = filename or self.tag
         path = Path(path).absolute()
         path.mkdir(parents=True, exist_ok=True)
-        name = f'{name}.png' if not name.endswith('.png') else name
-        name = get_available_file_name(str(path), name)
+        if not name.lower().endswith('.png'):
+            name = f'{name}.png'
 
         # 等待元素加载完成
         if self.tag == 'img':
@@ -423,7 +423,7 @@ class DriverElement(DrissionElement):
             while not self.run_script(js):
                 pass
 
-        img_path = f'{path}{sep}{name}'
+        img_path = str(get_usable_path(f'{path}{sep}{name}'))
         self.inner_ele.screenshot(img_path)
 
         return img_path

DrissionPage/driver_page.py

@@ -17,7 +17,7 @@ from selenium.webdriver.remote.webelement import WebElement
 from selenium.webdriver.support.wait import WebDriverWait
 
 from .base import BasePage
-from .common import get_available_file_name, format_html
+from .common import get_usable_path, format_html
 from .driver_element import DriverElement, make_driver_ele, _wait_ele
 from .session_element import make_session_ele
 
@@ -382,10 +382,11 @@ class DriverPage(BasePage):
         :return: 图片完整路径
         """
         name = filename or self.title
+        if not name.lower().endswith('.png'):
+            name = f'{name}.png'
         path = Path(path).absolute()
         path.mkdir(parents=True, exist_ok=True)
-        name = get_available_file_name(str(path), f'{name}.png')
-        img_path = f'{path}{sep}{name}'
+        img_path = str(get_usable_path(f'{path}{sep}{name}'))
         self.driver.save_screenshot(img_path)
         return img_path