如何在Python同一入口文件中为不同实验类配置独立argparse参数

解决方案 #

原代码的核心问题是第一次调用parse_args()时，若用户输入-h或参数不符合当前配置会直接终止程序，无法执行后续添加实验专属参数的逻辑，因此需要用分步解析的方式实现需求。

修改后代码 #

main.py

import argparse
import sys
# 导入各实验类后组装EXPERIMENTS字典
from experimentA import ExperimentA
from experimentB import ExperimentB

EXPERIMENTS = {
    "ExperimentA": ExperimentA,
    "ExperimentB": ExperimentB
}

def main(experiment):
    experiment()

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    # 新增公共参数可直接在此处添加
    parser.add_argument("experiment_name",
                        metavar='experiment_name',
                        type=str,
                        choices=EXPERIMENTS.keys(),
                        help="选择要运行的实验，可选值："+str(list(EXPERIMENTS.keys())),
                        )
    # 第一步：仅解析已知参数，剩余参数暂存，避免直接报错退出
    args, remaining_argv = parser.parse_known_args()
    # 获取选中的实验类，添加专属参数
    exp_cls = EXPERIMENTS[args.experiment_name]
    exp_cls.add_args(parser)
    # 第二步：全量解析所有参数
    args = parser.parse_args()
    # 初始化实验类
    experiment = exp_cls(args)
    
    sys.exit(main(experiment) or 0)

experimentA.py

class ExperimentA(Experiment):
    @classmethod
    def add_args(cls, parser):
        # 仅添加参数，不执行解析
        parser.add_argument('--expA-arg', help="实验A专属参数")
    
    def __init__(self, args):
        super().__init__()
        self.args = args
        self.param = self.args.expA_arg

    def __call__(self):
        # 实验A执行逻辑
        run_the_experiment_A()

experimentB.py

class ExperimentB(Experiment):
    @classmethod
    def add_args(cls, parser):
        parser.add_argument('--expB-arg', help="实验B专属参数")
    
    def __init__(self, args):
        super().__init__()
        self.args = args
        self.param = self.args.expB_arg

    def __call__(self):
        # 实验B执行逻辑
        run_the_experiment_B()

效果验证 #

完全符合预期的三个场景：

• 执行python3 main.py -h：仅展示公共参数和实验选择说明，无实验专属参数
• 执行python3 main.py ExperimentA -h：展示公共参数+实验A专属参数说明
• 执行python3 main.py ExperimentB -h：展示公共参数+实验B专属参数说明

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