Add init function to call build_pop as a module

build_pop.py

@@ -317,17 +317,10 @@ def confirm(deletions: dict, variants: dict):
     return input("Continue [Y/n]? ") in ["y", "Y", ""]
 
 
-def main():
-    args = parse_args()
-    reference = args.reference
-    sv_list = args.sv_list
-    output_dir = args.output_directory
-    tmp_dir = tempfile.mkdtemp()
-    haploid = args.haploid
-    nb_inds = args.nb_inds
-
+def init(output_dir, erase_output, sv_list, nb_inds, reference, proba_del, haploid, force_polymorphism,
+         coverage, read_len, insert_len_mean, insert_len_sd, threads, quiet=True):
     if os.path.isdir(output_dir):
-        if args.erase_output:
+        if erase_output:
             shutil.rmtree(output_dir)
         else:
             eprint("Error: output directory {0} already exists.".format(output_dir))
@@ -342,6 +335,8 @@ def main():
     if not os.path.isfile(reference + ".fai"):
         os.system("samtools faidx " + reference)
 
+    tmp_dir = tempfile.mkdtemp()
+
     prg_path = os.path.dirname(os.path.realpath(__file__))
 
     os.mkdir(output_dir)
@@ -357,21 +352,42 @@ def main():
 
     print("GENERATE RANDOM DELETIONS VARIANTS...\n")
     try:
-        sv_sim = VariantsSimulator(sv_list, args.threads)
-        deletions = sv_sim.get_random_deletions(args.proba_del, args.reference)
+        sv_sim = VariantsSimulator(sv_list, threads)
+        deletions = sv_sim.get_random_deletions(proba_del, reference)
         print("")
     except InputException as e:
         print(e)
         return False
-    if args.quiet or confirm(deletions, sv_sim.variants):
-        genotypes_for_inds = build_genotypes_vcf_list(deletions, output_vcf, haploid, args.force_polymorphism, nb_inds,
+    if quiet or confirm(deletions, sv_sim.variants):
+        genotypes_for_inds = build_genotypes_vcf_list(deletions, output_vcf, haploid, force_polymorphism, nb_inds,
                                                       tmp_dir, prg_path)
         build_fastas_chromosomes(reference, genotypes_for_inds, haploid, output_dir)
-        generate_samples_fastq(haploid, nb_inds, output_dir, args.coverage, args.read_len, args.insert_len_mean,
-                               args.insert_len_sd, prg_path, args.threads)
+        generate_samples_fastq(haploid, nb_inds, output_dir, coverage, read_len, insert_len_mean,
+                               insert_len_sd, prg_path, threads)
         print("DONE!\n")
     else:
         print("Aborted!\n")
 
+
+def main():
+    args = parse_args()
+    reference = args.reference
+    sv_list = args.sv_list
+    output_dir = args.output_directory
+    haploid = args.haploid
+    nb_inds = args.nb_inds
+    erase_output = args.erase_output
+    proba_del = args.proba_del
+    threads = args.threads
+    force_polymorphism = args.force_polymorphism
+    coverage = args.coverage
+    read_len = args.read_len
+    insert_len_mean = args.insert_len_mean
+    insert_len_sd = args.insert_len_sd
+    quiet = args.quiet
+    
+    init(output_dir, erase_output, sv_list, nb_inds, reference, proba_del, haploid, force_polymorphism,
+         coverage, read_len, insert_len_mean, insert_len_sd, threads, quiet)
+
 if __name__ == '__main__':
     sys.exit(main())