选文件时,只支持图片文件,防止发送别的资料

This commit is contained in:
liuwei
2025-02-24 16:36:40 +08:00
parent 2196037bd3
commit c64b58540d

View File

@@ -5,14 +5,19 @@ import random
def get_random_file_from_dir(directory): def get_random_file_from_dir(directory):
selected_file = None selected_file = None
file_count = 0 file_count = 0
# 定义图片文件的扩展名列表
image_extensions = {'.jpg', '.jpeg', '.png', '.gif', '.bmp', '.tiff', '.webp'}
# 遍历目录及其所有子目录 # 遍历目录及其所有子目录
for root, dirs, files in os.walk(directory): for root, dirs, files in os.walk(directory):
for file in files: for file in files:
# 计算当前文件被选中的概率 # 获取文件扩展名
_, ext = os.path.splitext(file)
# 只考虑图片文件
if ext.lower() in image_extensions:
file_count += 1 file_count += 1
if random.randint(1, file_count) == 1: if random.randint(1, file_count) == 1:
# 每个文件有1/file_count的概率被选中 # 每个图片文件有1/file_count的概率被选中
selected_file = os.path.join(root, file) selected_file = os.path.join(root, file)
return selected_file return selected_file
@@ -22,4 +27,7 @@ def get_xiuren_pic():
# 使用示例 # 使用示例
directory = '.' # 替换为你的目录路径 directory = '.' # 替换为你的目录路径
random_file_path = get_random_file_from_dir(directory) random_file_path = get_random_file_from_dir(directory)
if random_file_path:
return os.path.abspath(random_file_path) return os.path.abspath(random_file_path)
else:
return None