From f177fa7797d79f09806e3bc822adc1cdaca9a6d0 Mon Sep 17 00:00:00 2001
From: Grzegorz Latosinski <glatosinski@antmicro.com>
Date: Mon, 16 Nov 2020 17:38:07 +0100
Subject: [PATCH] Wrapped script in main function

Signed-off-by: Grzegorz Latosinski <glatosinski@antmicro.com>
---
 .../fet_simulator/fet_simulator.py            | 20 +++++++++++++------
 .../fet_simulator/generate-docs-plots.py      | 13 +++++++++---
 2 files changed, 24 insertions(+), 9 deletions(-)

diff --git a/scripts/python-skywater-pdk/skywater_pdk/fet_simulator/fet_simulator.py b/scripts/python-skywater-pdk/skywater_pdk/fet_simulator/fet_simulator.py
index 853a825..5af84a2 100644
--- a/scripts/python-skywater-pdk/skywater_pdk/fet_simulator/fet_simulator.py
+++ b/scripts/python-skywater-pdk/skywater_pdk/fet_simulator/fet_simulator.py
@@ -24,6 +24,7 @@ import matplotlib.pyplot as plt
 from pathlib import Path
 import csv
 from collections import defaultdict
+import sys
 
 logger = Logging.setup_logging()
 
@@ -120,14 +121,14 @@ def generate_fet_plots(
     # fet_W and fet_L values here are only for initialization, they are
     # later changed in the for loop
     c = create_test_circuit(fet_type, iparam, 0.15, 1, corner_path)
-    
+
     bins = read_bins(bins_csv)
 
     bins_by_W = defaultdict(list)
     # group bins by W
     for line in bins:
         bins_by_W[line[2]].append(line)
-    
+
     Ws = only_W if only_W is not None else list(bins_by_W.keys())
 
     for W in Ws:
@@ -149,17 +150,18 @@ def generate_fet_plots(
             )
             fg.tight_layout()
             fg.savefig(
-                Path(outdir) / (outprefix + f'_{name}_W{str(W).replace(".", "_")}.{ext}'),
+                Path(outdir) / (
+                    outprefix + f'_{name}_W{str(W).replace(".", "_")}.{ext}'),
                 bbox_extra_artists=(lg,),
                 bbox_inches='tight'
             )
         plt.close('all')
 
 
-if __name__ == '__main__':
+def main(argv):
     import argparse
 
-    parser = argparse.ArgumentParser()
+    parser = argparse.ArgumentParser(prog=argv[0])
     parser.add_argument(
         'fet_type',
         help='FET type to simulate'
@@ -194,7 +196,7 @@ if __name__ == '__main__':
         help='The image extension to use for figures',
         default='svg'
     )
-    args = parser.parse_args()
+    args = parser.parse_args(argv[1:])
 
     if args.outprefix is None:
         args.outprefix = args.fet_type
@@ -207,3 +209,9 @@ if __name__ == '__main__':
         args.outprefix,
         args.only_w,
         args.ext)
+
+    return 0
+
+
+if __name__ == '__main__':
+    sys.exit(main(sys.argv))
diff --git a/scripts/python-skywater-pdk/skywater_pdk/fet_simulator/generate-docs-plots.py b/scripts/python-skywater-pdk/skywater_pdk/fet_simulator/generate-docs-plots.py
index 4636cce..feb830c 100644
--- a/scripts/python-skywater-pdk/skywater_pdk/fet_simulator/generate-docs-plots.py
+++ b/scripts/python-skywater-pdk/skywater_pdk/fet_simulator/generate-docs-plots.py
@@ -19,11 +19,13 @@
 
 import argparse
 from pathlib import Path
+import sys
 
 from fet_simulator import generate_fet_plots
 
-if __name__ == '__main__':
-    parser = argparse.ArgumentParser()
+
+def main(argv):
+    parser = argparse.ArgumentParser(prog=argv[0])
     parser.add_argument(
         'fd_pr_dir',
         help='Path to the particular version of the primitive library',
@@ -35,7 +37,7 @@ if __name__ == '__main__':
         type=Path
     )
 
-    args = parser.parse_args()
+    args = parser.parse_args(argv[1:])
 
     typicalcorner = args.fd_pr_dir / 'models/corners/tt.spice'
 
@@ -70,3 +72,8 @@ if __name__ == '__main__':
             f'sim_{fetname}_',
             onlyw
         )
+    return 0
+
+
+if __name__ == '__main__':
+    sys.exit(main(sys.argv))