Merge branch 'master' of github.com:FofanovLab/MTSv

Author: tfursten

mtsv/mtsv_prep/MTSv_prune.py

@@ -399,27 +399,27 @@ def arg_unwrappers(args, arguments=None):
 
 def oneclickjson(path):
     arguments = []
-    for fh in os.listdir(os.path.join(path, "artifacts/")):
+    for fh in os.listdir(os.path.join(path, "artifacts")):
         if fnmatch.fnmatch(fh, "*.fas"):
             arguments.append({})
             fh = fh.split(".")[0]
-            arguments[-1]['serialization-path'] = os.path.abspath( os.path.join(path, "artifacts/{0}.p".format(fh)))
+            arguments[-1]['serialization-path'] = os.path.abspath( os.path.join(path, "artifacts","{0}.p".format(fh)))
 
-            arguments[-1]['fasta-path'] = os.path.abspath(os.path.join(path, "artifacts/{0}.fas".format(fh)))
+            arguments[-1]['fasta-path'] = os.path.abspath(os.path.join(path, "artifacts", "{0}.fas".format(fh)))
 
             arguments[-1]['minimum-length'] = 0
 
             arguments[-1]['maximum-length'] = float('inf')
-            arguments[-1]['taxdump-path'] = os.path.abspath(os.path.join(path, "artifacts/taxdump.tar.gz"))
+            arguments[-1]['taxdump-path'] = os.path.abspath(os.path.join(path, "artifacts","taxdump.tar.gz"))
 
             arguments[-1]['acc-to-taxid-paths'] = []
-            for fp in os.listdir(os.path.join(path,"artifacts/")):
+            for fp in os.listdir(os.path.join(path,"artifacts")):
                     if fnmatch.fnmatch(fp, "*accession2taxid*"):
-                        arguments[-1]['acc-to-taxid-paths'].append(os.path.abspath(os.path.join(path,"artifacts/",fp)))
+                        arguments[-1]['acc-to-taxid-paths'].append(os.path.abspath(os.path.join(path,"artifacts",fp)))
 
             arguments[-1]['rollup-rank'] = "species"
 
-            with open(os.path.abspath(os.path.join(path,"artifacts/{0}.json".format(fh))), "w") as file:
+            with open(os.path.abspath(os.path.join(path,"artifacts", "{0}.json".format(fh))), "w") as file:
                 json.dump(arguments[-1], file, sort_keys=True, indent=4)
 
     return arguments
@@ -474,7 +474,10 @@ def ftp_dl(x):
     connection.login()
 
     while x:
-        fp_path = x.pop()
+        try:
+            fp_path = x.pop()
+        except IndexError:
+            break
         raw_path = fp_path[0]
         fp_path = fp_path[1]
         try:
@@ -582,17 +585,18 @@ def pull(path="",thread_count=1,databases ={"genbank"} ):
                 continue
         except:
             pass
-        if line[11].strip().decode().lower().replace(" ", "_") in databases:
+        db = line[11].strip().decode().lower().replace(" ", "_")
+        if db in databases:
             try:
                 temp = line[19].split(ftp_path.encode(),1)[1].decode()
                 temp_path = "{0}/{1}_genomic.gbff.gz".format(temp, os.path.basename(temp))
             except:
                 continue
             try:
-                level2path[line[11].strip()].append(temp_path)
+                level2path[db.encode()].append(temp_path)
             except:
-                level2path[line[11].strip()] = [temp_path]
-            to_download.append((os.path.join(raw_path,"flat_files/"),temp_path))
+                level2path[db.encode()] = [temp_path]
+            to_download.append((os.path.join(raw_path,"flat_files"),temp_path))
 
     reader = BytesIO()
     connection.retrbinary("RETR {0}{1}".format(assembly_gb,assembly_gb_summary) ,reader.write)
@@ -609,32 +613,32 @@ def pull(path="",thread_count=1,databases ={"genbank"} ):
                 continue
         except:
             pass
-
-        if line[11].strip().decode().lower().replace(" ", "_") in databases:
+        db = line[11].strip().decode().lower().replace(" ", "_")
+        if db in databases:
             try:
                 temp = line[19].split(ftp_path.encode(),1)[1].decode()
                 temp_path = "{0}/{1}_genomic.gbff.gz".format(temp, os.path.basename(temp))
             except:
                 continue
             try:
-                level2path[line[11].strip()].append(temp_path)
+                level2path[db.encode()].append(temp_path)
             except:
-                level2path[line[11].strip()] = [temp_path]
+                level2path[db.encode()] = [temp_path]
 
             to_download.append((os.path.join(raw_path,"flat_files/"), temp_path) )
     artifacts = [(os.path.join(raw_path,"artifacts/"),"/pub/taxonomy/taxdump.tar.gz")]
     tax_path = "/pub/taxonomy/accession2taxid/"
     for file in connection.nlst(tax_path):
         if not fnmatch.fnmatch(os.path.basename(file), 'dead*') and not fnmatch.fnmatch(file, '*md5'):
-            artifacts.append((os.path.join(raw_path,"artifacts/"),file))
+            artifacts.append((os.path.join(raw_path, "artifacts"), file))
 
     man = Manager()
     connection.quit()
     to_download += artifacts
     to_download = man.list(to_download)
 
     proc = [Process(target = ftp_dl, args=(to_download,)) for i in range(thread_count)]
-    with open(os.path.join(raw_path,"artifacts/ftp_dl.log"), "w"):
+    with open(os.path.join(raw_path,"artifacts", "ftp_dl.log"), "w"):
         pass
     for p in proc:
         p.start()
@@ -647,7 +651,7 @@ def pull(path="",thread_count=1,databases ={"genbank"} ):
         fp = "{0}_ff.txt".format(i.decode().replace(" ","_"))
         with open(os.path.join(raw_path,"artifacts/",fp), "w") as out_file:
             for line in level2path[i]:
-                out_file.write("{0}\n".format(os.path.join(os.path.abspath(raw_path),"flat_files/",os.path.basename(line))))
+                out_file.write("{0}\n".format(os.path.join(os.path.abspath(raw_path),"flat_files/", os.path.basename(line))))
     return string_date
 
 if __name__ =="__main__":

mtsv/mtsv_prep/main.py

@@ -77,12 +77,12 @@ def decompression(x, path):
     return x.strip(".gz")
 
 def oneclickbuild(args):
-    with open(os.path.join(args.path, "artifacts/decompression.log"), "w" ):
+    with open(os.path.join(args.path, "artifacts","decompression.log"), "w" ):
         pass
     pool = Pool(args.threads)
     pool.starmap(decompression, [(os.path.abspath(x),
-                                  args.path) for x in iglob(os.path.join(args.path,"flat_files/*.gz"))])
-    for fp in iglob(os.path.join(args.path,"artifacts/","*_ff.txt")):
+                                  args.path) for x in iglob(os.path.join(args.path,"flat_files", "*.gz"))])
+    for fp in iglob(os.path.join(args.path,"artifacts","*_ff.txt")):
         db = list(os.path.split(fp))
         db[1] = db[1].strip().replace("_ff.txt",".fas")
         db = os.path.abspath("{0}{1}{2}".format(db[0],os.sep,db[1]))
@@ -100,14 +100,18 @@ def oneclickbuild(args):
         build_db(os.path.abspath(fp), db, os.devnull, os.devnull, args.threads, os.devnull)
 
     arguments = oneclickjson(args.path)
-
+    for argument in arguments:
+        pool.apply_async(tree_make, (argument['taxdump-path'], ))
+        break
     pool.starmap(acc_serialization, [(argument['acc-to-taxid-paths'], argument['fasta-path'],
                                       argument['taxdump-path']) for argument in arguments ])
 
+    pool.close()
     # shutil.rmtree(os.path.join(args.path, "flat_files" ))
 
-def mapper(x):
-    return clip(*x)
+def tree_make(in_path):
+    out_path = os.path.abspath(os.path.join(os.path.dirname(in_path), "tree.index"))
+    subprocess.run("{0} --index {1} --dump {2}".format(bin_path('mtsv-tree-build'), out_path, os.path.abspath(in_path) ).split())
 
 def partition(args):
     partition_list = []
@@ -247,6 +251,14 @@ def json_updater(args):
             params['partition-path'].append(os.path.abspath(path))
 
         params['tree-index'] = os.path.abspath(os.path.join(args.path, "artifacts","tree.index"))
+        params['serialization-path'] = os.path.abspath(os.path.join(args.path, "artifacts","{0}.p".format(base)))
+        params['taxdump-path'] = os.path.abspath(os.path.join(args.path, "artifacts","taxdump.tar.gz"))
+        params['fasta-path'] = os.path.abspath(os.path.join(args.path, "artifacts","{0}.fas".format(base)))
+
+        params['acc-to-taxid-paths'] = []
+        for path in iglob(os.path.join(args.path, "artifacts","*taxid.gz")):
+            params['acc-to-taxid-paths'].append(os.path.abspath(path))
+
         with open(json_path, "w") as file:
             json.dump(params, file, sort_keys=True, indent=4)
 
@@ -340,43 +352,45 @@ def setup_and_run(parser):
         args = parser.parse_known_args()[0]
 
         try:
-
             if args.cmd_class == Database:
+                for i, val in enumerate(args.includedb):
+                    args.includedb[i] = val.strip().replace(" ","_").lower()
                 if args.download_only:
                     args.path = os.path.abspath(oneclickdl(args))
-
                 elif args.build_only:
                     if args.path and  os.path.isdir(args.path):
                         oneclickbuild(args)
                         json_updater(args)
                         make_json_abs(args)
-
                     else:
-                        print("A valid path was not specified")
+                        print("A valid directory path was not specified")
                 else:
                     args.path = os.path.abspath(oneclickdl(args))
                     oneclickbuild(args)
                     json_updater(args)
                     make_json_abs(args)
-
             elif args.cmd_class == CustomDB:
+                for i, val in enumerate(args.customdb):
+                    args.includedb[i] = val.strip().replace(" ","_").lower()
                 oneclickfmbuild(args, args.partitions == DEFAULT_PARTITIONS)
                 json_updater(args)
                 make_json_abs(args)
 
-            # try:
-            #     make_json_rel(args)
-            # except:
-            #     pass
 
         except AttributeError:
             sys.argv[1] = "database"
             args = parser.parse_known_args()[0]
+            for i, val in enumerate(args.includedb):
+                args.includedb[i] = val.strip().replace(" ", "_").lower()
             args.path = os.path.abspath(oneclickdl(args))
             oneclickbuild(args)
             path = args.path
+            json_updater(args)
+            make_json_abs(args)
             sys.argv[1] = "custom_db"
             args = parser.parse_known_args()[0]
+            for i, val in enumerate(args.customdb):
+                args.customdb[i] = val.strip().replace(" ", "_").lower()
             args.path = path
             oneclickfmbuild(args, args.partitions == DEFAULT_PARTITIONS)
             json_updater(args)