使用argparse传参时Python解压脚本抛出StopIteration错误求助

问题现象 #

脚本在无参数使用默认配置时运行正常，但传入-z指定目录或-f指定ZIP文件路径时，抛出StopIteration错误，错误栈指向unzip_templates函数中查找ZIP文件的代码段：

235: An error occurred:
Traceback (most recent call last):
  File "/Users/ops.eekfonky/code/dev/server/trunk/deployment/./templatesToPR.py", line 217, in main
    template_name, filename = unzip_templates(zip_dir, extraction_dir)
                              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ops.eekfonky/code/dev/server/trunk/deployment/./templatesToPR.py", line 57, in unzip_templates
    filename = next(f for f in os.listdir(zip_dir) if f.endswith('.zip'))
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
StopIteration

错误原因 #

1. 参数逻辑冲突 #

argparse中-z参数设置了默认值，导致即使传入-f参数，args.zip_dir依然保留默认值，使得main函数中的if args.zip_dir:分支始终优先执行，不会走到elif args.zip_file的处理逻辑。脚本会错误地将默认目录传入unzip_templates，如果该目录下没有ZIP文件，就触发next()的StopIteration异常。

2. next()的不安全使用 #

直接用next()遍历生成器时，若没有匹配的ZIP文件，生成器会耗尽，直接抛出StopIteration，没有任何异常捕获或 fallback 处理。

修复方案 #

1. 修复参数解析与分支判断 #

首先修改参数解析逻辑，让-z和-f互斥（同一时间只能用一个），避免逻辑冲突：

def parse_arguments():
    """Parses command-line arguments."""
    parser = argparse.ArgumentParser(description='''
    This script handles git repositories and templates based on internal defaults or command-line overrides.
    Use command-line options to specify custom settings for the script's operation.''')
    parser.add_argument('--def-branch', '-d', default=DEFAULTS["default_branch"],
                        help='Default git branch name. Default is "%(default)s".')
    # 创建互斥组，确保-z和-f只能二选一
    group = parser.add_mutually_exclusive_group()
    group.add_argument('--zip-dir', '-z', default=DEFAULTS["zip_dir"],
                        help='Directory for zipped templates. Default is "%(default)s".')
    group.add_argument('--zip-file', '-f', type=str,
                        help='Specify the .zip filename explicitly.')
    return parser.parse_args()

然后修改main中的路径判断逻辑，优先处理-f参数：

def main():
    args = parse_arguments()
    install_and_import('GitPython', 'git')
    from git import Repo, GitCommandError

    script_dir = os.path.dirname(os.path.abspath(__file__))
    os.chdir(script_dir)
    git_repo_root = os.path.abspath(os.path.join(script_dir, "../../../"))
    template_dir = os.path.join(git_repo_root, "server/trunk/services/src/main/resources/templates")

    # 优先处理-f指定的单个ZIP文件
    if args.zip_file:
        zip_dir = os.path.abspath(os.path.expanduser(args.zip_file))
        if not os.path.isfile(zip_dir) or not zip_dir.lower().endswith('.zip'):
            raise ValueError("Error: When using the -f flag, you must provide a valid .zip file path.")
    else:
        # 处理默认或-z指定的目录
        zip_dir = os.path.abspath(os.path.expanduser(args.zip_dir))
        if not os.path.isdir(zip_dir):
            raise ValueError("Error: When using the -z flag, you must provide a directory containing ZIP files.")
        ensure_directory_exists(zip_dir)
    
    # ... 后续原有代码不变

2. 安全获取ZIP文件 #

修改unzip_templates函数，先收集所有匹配的ZIP文件，再进行判断，避免StopIteration：

def unzip_templates(zip_dir, extraction_path):
    """Extracts a template ZIP, handling 'email' and 'sms' directories and preserving nested structures."""
    zip_file_path = None

    if os.path.isdir(zip_dir):
        # 收集目录下所有ZIP文件（不区分大小写）
        zip_files = [f for f in os.listdir(zip_dir) if f.lower().endswith('.zip')]
        if not zip_files:
            raise FileNotFoundError(f"No .zip files found in directory: {zip_dir}")
        # 取第一个ZIP文件
        filename = zip_files[0]
        zip_file_path = os.path.join(zip_dir, filename)
    else:
        filename = os.path.basename(zip_dir)
        zip_file_path = zip_dir
        # 验证传入的文件是ZIP格式
        if not filename.lower().endswith('.zip'):
            raise ValueError(f"Provided file is not a .zip file: {zip_file_path}")

    # ... 后续原有代码不变

3. 优化cleanup函数的路径逻辑 #

修复cleanup中默认路径的问题，确保删除正确的ZIP文件：

def cleanup(extraction_dir, args, filename=None):
    """Clean up temporary files and optionally the original ZIP file(s)."""

    shutil.rmtree(extraction_dir)
    logging.info(f"Temporary files cleaned up from {extraction_dir}")

    # 根据用户参数确定要删除的ZIP文件路径
    if args.zip_file:
        original_zip_file = os.path.abspath(os.path.expanduser(args.zip_file))
    else:
        zip_dir = os.path.abspath(os.path.expanduser(args.zip_dir))
        original_zip_file = os.path.join(zip_dir, filename)

    if os.path.isfile(original_zip_file):
        os.remove(original_zip_file)
        logging.debug(f"Original ZIP file removed: {original_zip_file}")
    else:
        logging.warning(f"ZIP file not found for deletion: {original_zip_file}")

内容的提问来源于stack exchange，提问作者eekfonky