如何通过argparse覆盖输出路径 将输出文件保存到指定目录

我有如下示例代码，使用argparse实现代码的调用与执行。
test1.py代码如下：

#test1.py
import numpy
from in_para import *
from matplotlib import pyplot as plt

nx = 31
dx = 2 / (nx - 1)
nt = 20   
D = 0.3  
sigma = .2 
dt = sigma * dx**2 / D 

u = numpy.ones(nx)      
u[int(.5 / dx):int(1 / dx + 1)] = 2  

un = numpy.ones(nx) 

for n in range(nt):  
    un = u.copy() 
    for i in range(1, nx - 1):
        u[i] = un[i] + D * dt / dx**2 * (un[i+1] - 2 * un[i] + un[i-1])
        
def test1():
    for i in range(nx):
        print(i,u[i])
        f = open('Output.txt', 'w')
        print(i,u[i], file = f)
    plt.plot(numpy.linspace(0, 2, nx), u)
    plt.savefig('test1.png', bbox_inches='tight')
    plt.show()


if __name__ == '__main__':
    test1()

通过argparse调用上述代码的main.py脚本如下：

import numpy
from test1 import *
import argparse
import sys

from matplotlib import pyplot as plt


def main():
    parser = argparse.ArgumentParser() 
    parser.add_argument('-C','--Modul',type=str, help='Choose of the component', default=None)
    args = vars(parser.parse_args())
 
    if args['Modul'] == 'test':
        print("This is output for test1:")
        test1()
   
          
if __name__=='__main__':
    main()

我目前正在开发Python流水线，希望通过命令行接口（CLI）传入argparse参数的方式覆盖默认输出文件路径。test1.py仅为示例代码，实际业务场景中代码体量较大，我不希望main.py的使用者手动进入模块内部修改代码。需要实现通过argparse覆盖输出文件路径，收集test1.py的输出内容，将.txt和.png文件写入指定目录、使用自定义文件名，预期的命令行调用格式如下：

python3 main.py -C test -O [Output options] 

完全可以实现，只需要做两处低侵入修改，不需要改动核心计算逻辑：

• 给模块入口函数增加可选的输出路径参数，保留默认值，完全兼容原有调用方式
• 在CLI入口的argparse配置中新增输出参数，解析后把路径传入对应模块函数即可，同时自动处理目录不存在的异常场景

第一步：修改test1.py #

只需要调整test1函数的定义和文件写入逻辑，其余计算代码不需要动。顺便修复原代码里的写文件bug：原逻辑在循环内每次用w模式打开文件，会清空之前写入的内容，最终txt文件只会保留最后一行数据。

import os
# 保留原有所有导入和顶部的计算逻辑不变，只修改test1函数
def test1(output_txt_path="Output.txt", output_png_path="test1.png"):
    # 自动创建输出文件所在目录，避免目录不存在时报错
    os.makedirs(os.path.dirname(os.path.abspath(output_txt_path)), exist_ok=True)
    os.makedirs(os.path.dirname(os.path.abspath(output_png_path)), exist_ok=True)

    # 把文件打开操作移到循环外，用上下文管理器自动处理文件关闭
    with open(output_txt_path, 'w') as f:
        for i in range(nx):
            print(i, u[i])
            print(i, u[i], file=f)
    plt.plot(numpy.linspace(0, 2, nx), u)
    plt.savefig(output_png_path, bbox_inches='tight')
    plt.show()

第二步：修改main.py #

新增-O/--output参数，支持传入输出路径前缀，自动拼接txt和png后缀，不传参时自动走默认路径逻辑：

import os
import numpy
from test1 import test1
import argparse
from matplotlib import pyplot as plt


def main():
    parser = argparse.ArgumentParser() 
    parser.add_argument('-C','--Modul',type=str, help='Choose of the component', default=None)
    parser.add_argument('-O','--output',type=str, help='Output path prefix, will generate {prefix}.txt and {prefix}.png', default=None)
    args = parser.parse_args()
 
    if args.Modul == 'test':
        print("This is output for test1:")
        if args.output:
            txt_path = f"{args.output}.txt"
            png_path = f"{args.output}.png"
            test1(output_txt_path=txt_path, output_png_path=png_path)
        else:
            test1()
   
          
if __name__=='__main__':
    main()

调用方式 #

修改完成后即可按预期格式使用，比如要把结果输出到./result/exp2目录下，文件命名为diffusion_result.txt和diffusion_result.png，执行以下命令即可：

python3 main.py -C test -O ./result/exp2/diffusion_result

如果需要支持其他模块的自定义输出，用同样的方式给对应模块的入口函数加可选路径参数即可，不需要使用者深入模块内部修改硬编码路径，整体改动量极小。

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