Python中筛选特定命名规则的EXR文件方案优化问询

我正在设计一款工具，需从输入文件夹中仅提取符合以下命名规则的EXR图像文件：u#_v#.exr 或 u#_v#_#.exr（其中#表示正非零整数），其余文件均需忽略。以下是我当前可运行的代码，请问是否存在更优或更高效的实现方式？

def main():
    # Add and read command line arguments
    parser = argparse.ArgumentParser()
    parser.add_argument('--input_folder', type=str, help='Directory where input images are located')
    parser.add_argument('--output_folder', type=str, help='Directory where output image should be written')
    args = parser.parse_args()

    # Change directory to input folder and check all filenames belonging to our convention
    os.chdir(args.input_folder)
    all_files = check_combinatons_of_numeric_characters('u_v.exr', 'u_v_.exr')
    print(all_files)


def check_combinatons_of_numeric_characters(convention1, convention2):
    # Combinations for first convention which was supplied
    split_convention1 = convention1.split('_')
    convention1_combinations_alpha = np.array([])

    convention1_combination1 = glob.glob(split_convention1[0] + '[0-9]_' +
                                         split_convention1[1].split('.')[0] + '[0-9].' +
                                         split_convention1[1].split('.')[1]
                                         )

    convention1_combination2 = glob.glob(split_convention1[0] + '[0-9][0-9]_' +
                                         split_convention1[1].split('.')[0] + '[0-9][0-9].' +
                                         split_convention1[1].split('.')[1]
                                         )

    convention1_combination3 = glob.glob(split_convention1[0] + '[0-9][0-9]_' +
                                         split_convention1[1].split('.')[0] + '[0-9].' +
                                         split_convention1[1].split('.')[1]
                                         )

    convention1_combination4 = glob.glob(split_convention1[0] + '[0-9]_' +
                                         split_convention1[1].split('.')[0] + '[0-9][0-9].' +
                                         split_convention1[1].split('.')[1]
                                         )

    convention1_combinations_alpha = np.concatenate((convention1_combination1,
                                                     convention1_combination2,
                                                     convention1_combination3,
                                                     convention1_combination4),
                                                    )

    # Combinations for second convention supplied
    split_convention2 = convention2.split('_')
    convention2_combinations_alpha = np.array([])

    convention2_combination1 = glob.glob(split_convention2[0] + '[0-9]_'+
                                         split_convention2[1] + '[0-9]_[0-9]' +
                                         split_convention2[2]
                                         )
    convention2_combination2 = glob.glob(split_convention2[0] + '[0-9][0-9]_' +
                                         split_convention2[1] + '[0-9]_[0-9]' +
                                         split_convention2[2]
                                         )
    convention2_combination3 = glob.glob(split_convention2[0] + '[0-9]_' +
                                         split_convention2[1] + '[0-9]_[0-9][0-9]' +
                                         split_convention2[2]
                                         )
    convention2_combination4 = glob.glob(split_convention2[0] + '[0-9][0-9]_' +
                                         split_convention2[1] + '[0-9][0-9]_[0-9]' +
                                         split_convention2[2]
                                         )
    convention2_combination5 = glob.glob(split_convention2[0] + '[0-9]_' +
                                         split_convention2[1] + '[0-9][0-9]_[0-9][0-9]' +
                                         split_convention2[2]
                                         )
    convention2_combination6 = glob.glob(split_convention2[0] + '[0-9][0-9]_' +
                                         split_convention2[1] + '[0-9]_[0-9][0-9]' +
                                         split_convention2[2]
                                         )

    convention2_combinations_alpha = np.concatenate((convention2_combination1,
                                                     convention2_combination2,
                                                     convention2_combination3,
                                                     convention2_combination4,
                                                     convention2_combination5,
                                                     convention2_combination6),
                                                    )

    list_of_files = np.concatenate((convention1_combinations_alpha, convention2_combinations_alpha))

    return list_of_files

if __name__ == '__main__':
    main()

优化思路与实现 #

原代码通过枚举1位、2位数字的组合来匹配文件名，不仅冗余，还无法支持3位及以上的正整数。下面提供两种更简洁高效的优化方案：

方案1：正则表达式直接匹配（最灵活） #

用正则表达式精准匹配任意长度的正非零整数，同时避免切换工作目录的风险。

import argparse
import os
import re

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--input_folder', type=str, help='输入图像所在目录')
    parser.add_argument('--output_folder', type=str, help='输出图像保存目录')
    args = parser.parse_args()

    # 正则规则：匹配u+正整数_v+正整数(_+正整数)?.exr，支持大小写后缀
    pattern = re.compile(r'^u[1-9]\d*_v[1-9]\d*(?:_[1-9]\d*)?\.exr$', re.IGNORECASE)
    
    all_files = []
    # 遍历文件夹，筛选符合规则的文件
    for filename in os.listdir(args.input_folder):
        if pattern.match(filename):
            full_path = os.path.join(args.input_folder, filename)
            if os.path.isfile(full_path):
                all_files.append(full_path)
    
    print(all_files)

if __name__ == '__main__':
    main()

方案2：Glob粗筛+正则精筛（更高效） #

先用glob过滤出大致符合前缀的文件，再用正则二次筛选，减少需要匹配的文件数量。

import argparse
import os
import re
import glob

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--input_folder', type=str, help='输入图像所在目录')
    parser.add_argument('--output_folder', type=str, help='输出图像保存目录')
    args = parser.parse_args()

    pattern = re.compile(r'^u[1-9]\d*_v[1-9]\d*(?:_[1-9]\d*)?\.exr$', re.IGNORECASE)
    # 先用glob粗筛，缩小匹配范围
    candidate_files = glob.glob(os.path.join(args.input_folder, 'u*_v*.exr'))
    candidate_files += glob.glob(os.path.join(args.input_folder, 'u*_v*_*.exr'))
    
    # 正则精筛
    all_files = [f for f in candidate_files if pattern.match(os.path.basename(f))]
    
    print(all_files)

if __name__ == '__main__':
    main()

优化点说明 #

1. 正则表达式：支持任意长度的正非零整数，无需枚举数字位数组合，覆盖所有符合规则的文件名
2. 路径处理：直接拼接完整路径，避免os.chdir改变全局工作目录带来的潜在问题
3. 效率提升：方案2通过glob先过滤掉大部分无关文件，减少正则匹配的工作量
4. 兼容性：添加re.IGNORECASE，支持.EXR等大写后缀的文件

内容的提问来源于stack exchange，提问作者Ahmad Uzzam Masood