How to override output filepath and write it along with output files at a desired directory through argparse

I have the following example code and I use argparse to call and execute it.

#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()

Below is the main.py script where I call the above code via argparse

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()

Since I’m working on a python pipeline, I would like to override the output filepath and give it via CLI, i.e., as an arseparse argument. As I’ve said in the beginning, the test1.py is an example code, in reality it will have several lines of code and I dont want the user of main.py to go inside and modify it manually. Thus, my question is, can the output filepath be overridden via argparse ? such that I can collect the outputs from test1.py and write the .txt and .png file at a desired directory and desired filenames. For example, the command line should look like this:

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

Thanks

Answer

You can do it adding this line to main.py:

import os
parser.add_argument('-O','--Output',type=str, help='Output Directory', default="./Output")
output_dir = args['Output']
if not os.path.exists(output_dir):
    os.makedirs(output_dir)

After that pass output directory to test function:

test1(output_dir)

Change the test function to accept the output directory and change the path in plt.savefig

def test1(output_dir):
    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(os.path.join(output_dir, 'test1.png'), bbox_inches='tight')
    plt.show()