change agen main to be more conventional

mingkaic · Dec 10, 2018 · 73ca9fe · 73ca9fe
1 parent 476c240
commit 73ca9fe
Showing 1 changed file with 28 additions and 30 deletions.
diff --git a/age/agen.py b/age/agen.py
@@ -189,15 +189,37 @@ def make_dir(fields, includes, includepath, gen_capi):
 
     return out
 
-def main(cfgpath = None,
-    outpath = None,
-    strip_prefix = '',
-    gen_capi = False):
+def str2bool(opt):
+    optstr = opt.lower()
+    if optstr in ('yes', 'true', 't', 'y', '1'):
+        return True
+    elif optstr in ('no', 'false', 'f', 'n', '0'):
+        return False
+    else:
+        raise argparse.ArgumentTypeError('Boolean value expected.')
+
+def main(args):
+
+    parser = argparse.ArgumentParser(description=prog_description)
+    parser.add_argument('--cfg', dest='cfgpath', nargs='?',
+        help='Configuration json file on mapping info (default: read from stdin)')
+    parser.add_argument('--out', dest='outpath', nargs='?',
+        help='Directory path to dump output files (default: write to stdin)')
+    parser.add_argument('--strip_prefix', dest='strip_prefix', nargs='?', default='',
+        help='Directory path to dump output files (default: write to stdin)')
+    parser.add_argument('--gen_capi', dest='gen_capi',
+        type=str2bool, nargs='?', const=True, default=False,
+        help='Whether to generate C api or not (default: False)')
+    args = parser.parse_args(args)
+
+    outpath = args.outpath
+    strip_prefix = args.strip_prefix
 
     includepath = outpath
     if includepath and includepath.startswith(strip_prefix):
         includepath = includepath[len(strip_prefix):].strip("/")
 
+    cfgpath = args.cfgpath
     if cfgpath:
         with open(str(cfgpath), 'r') as cfg:
             cfg_str = cfg.read()
@@ -207,7 +229,7 @@ def main(cfgpath = None,
         cfg_str = sys.stdin.read()
     fields, includes = parse(cfg_str)
 
-    directory = make_dir(fields, includes, includepath, gen_capi)
+    directory = make_dir(fields, includes, includepath, args.gen_capi)
 
     if outpath:
         for fname, content in directory:
@@ -218,29 +240,5 @@ def main(cfgpath = None,
             print("============== %s ==============" % fname)
             print(content)
 
-def str2bool(opt):
-    optstr = opt.lower()
-    if optstr in ('yes', 'true', 't', 'y', '1'):
-        return True
-    elif optstr in ('no', 'false', 'f', 'n', '0'):
-        return False
-    else:
-        raise argparse.ArgumentTypeError('Boolean value expected.')
-
 if '__main__' == __name__:
-    parser = argparse.ArgumentParser(description=prog_description)
-    parser.add_argument('--cfg', dest='cfgpath', nargs='?',
-        help='Configuration json file on mapping info (default: read from stdin)')
-    parser.add_argument('--out', dest='outpath', nargs='?',
-        help='Directory path to dump output files (default: write to stdin)')
-    parser.add_argument('--strip_prefix', dest='strip_prefix', nargs='?', default='',
-        help='Directory path to dump output files (default: write to stdin)')
-    parser.add_argument('--gen_capi', dest='gen_capi',
-        type=str2bool, nargs='?', const=True, default=False,
-        help='Whether to generate C api or not (default: False)')
-    args = parser.parse_args()
-
-    main(cfgpath = args.cfgpath,
-        outpath = args.outpath,
-        strip_prefix = args.strip_prefix,
-        gen_capi = args.gen_capi)
+    main(sys.argv[1:])