diff --git a/.gitignore b/.gitignore index 83afd5e..b71d49e 100755 --- a/.gitignore +++ b/.gitignore @@ -7,8 +7,12 @@ prepUpdate.py *.pyc include.txt build.sh -update_mtdb -extract_mtdb -predb2mtdb -mtdb -manage_mtdb +# legacy root-level console-script copies (anchored so they do NOT clobber the +# mycotools/mtdb/ package or the test/update_mtdb/ suite of the same name) +/update_mtdb +/extract_mtdb +/predb2mtdb +/mtdb +/manage_mtdb +.gitignore +CLAUDE.md diff --git a/Dockerfile b/Dockerfile index dc3ef1e..c2f9bea 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,7 +1,6 @@ FROM mambaorg/micromamba:2.3.0-ubuntu22.04 AS app -ARG MYCOTOOLS_VER="1.0.0" -USER root +ARG MYCOTOOLS_VER="2.0.0" # 'LABEL' instructions tag the image with metadata that might be important to the user LABEL base.image="mambaorg/micromamba:2.3.0-ubuntu22.04" @@ -12,21 +11,28 @@ LABEL description="Mycotools is a compilation of computational biology tools and LABEL website="https://github.com/xonq/mycotools" LABEL license="https://github.com/xonq/mycotools/blob/master/LICENSE" LABEL maintainer="Zachary Konkel" -LABEL maintainer.email="konkelzach@protonmail.com" # this is unfortunately necessary to install ete4 RUN micromamba install --name base -c conda-forge -c bioconda -c defaults legacy-cgi pip mycotools=${MYCOTOOLS_VER} && \ eval "$(micromamba shell hook --shell bash)" && \ micromamba activate base && \ python3 -m pip install dna_features_viewer && \ - micromamba clean -a -f -y && \ - mkdir /data + micromamba clean -a -f -y + +USER root + +RUN mkdir /data + +RUN chown $MAMBA_USER:$MAMBA_USER /data + +USER MAMBA_USER ENV PATH="/opt/conda/bin/:${PATH}" \ LC_ALL=C.UTF-8 # 'CMD' instructions set a default command when the container is run. This is typically 'tool --help.' CMD [ "mtdb", "--help" ] +CMD [ "mycotools", "--help" ] # 'WORKDIR' sets working directory -WORKDIR /data \ No newline at end of file +WORKDIR /data diff --git a/mycotools/acc2fq.py b/mycotools/acc2fq.py deleted file mode 100755 index 1cd278f..0000000 --- a/mycotools/acc2fq.py +++ /dev/null @@ -1,220 +0,0 @@ -#! /usr/bin/env python3 - -import os -import re -import sys -import gzip -import argparse -from Bio import SeqIO -from collections import defaultdict - -# from mycotools.lib.biotools import dict2fq -from mycotools.lib.kontools import format_path, eprint, stdin2str - - -def dict2fq(fastq_dict, description=True): - """Convert a fastq dictionary to a fastq string""" - fastq_string = "" - if description: - for seq in fastq_dict: - fastq_string += "@" + seq.rstrip() - if "description" in fastq_dict[seq]: - fastq_string += (" " + fastq_dict[seq]["description"]).rstrip() + "\n" - else: - fastq_string += "\n" - fastq_string += ( - fastq_dict[seq]["sequence"].rstrip() - + "\n+\n" - + fastq_dict[seq]["score"].rstrip() - + "\n" - ) - - else: - for seq in fastq_dict2: - fastq_string += ( - "@" - + seq.rstrip() - + "\n" - + fastq_dict[seq]["sequence"].rstrip() - + "\n+\n" - + fastq_dict[seq]["score"].rstrip() - + "\n" - ) - - return fastq_string - - -def acc2fq(fq_path, accs): - fq_dict, parse, header, indices, count = {}, False, False, [], 0 - if fq_path.endswith((".gz", ".gzip")): - with gzip.open(fq_path, "rt") as raw: - for line in raw: - data = line.rstrip() - if not parse: - if not accs: - break - if data.startswith("@"): - header = data[1:].split(" ") - seq_name = header[0] - if seq_name in accs: - fq_dict[seq_name] = { - "sequence": "", - "description": " ".join(header[1:]), - "score": "", - } - parse = True - append_seq = "sequence" - accs.remove(seq_name) - elif parse: - if data == "+" and append_seq != "score": - append_seq = "score" - else: - fq_dict[seq_name][append_seq] += data - seq_len = len(fq_dict[seq_name]["sequence"]) - scr_len = len(fq_dict[seq_name]["score"]) - # the FASTQ is a terribly inefficient format - if seq_len <= scr_len: - if seq_len == scr_len: - parse = False - else: - raise ValueError( - "discrepancy between " - + "sequence and score " - + f"lengths: {seq_name}" - ) - - else: - with open(fq_path, "r") as raw: - for line in raw: - data = line.rstrip() - if not parse: - if not accs: - break - - if data.startswith("@"): - parse = False - header = data[1:].split(" ") - seq_name = header[0] - if seq_name in accs: - fq_dict[seq_name] = { - "sequence": "", - "description": " ".join(header[1:]), - "score": "", - } - parse = True - append_seq = "sequence" - accs.remove(seq_name) - elif parse: - if data == "+" and append_seq != "score": - append_seq = "score" - else: - fq_dict[seq_name][append_seq] += data - seq_len = len(fq_dict[seq_name]["sequence"]) - scr_len = len(fq_dict[seq_name]["score"]) - if seq_len <= scr_len: - if seq_len == scr_len: - parse = False - else: - raise ValueError( - "discrepancy between " - + "sequence and score " - + f"lengths: {seq_name}" - ) - - return fq_dict - - -def acc2fq_pe(fq_path, indices): - fq_dict, parse, count = {}, False, [] - if fq_path.endswith((".gz", ".gzip")): - with gzip.open(fq_path, "rt") as raw: - for line in raw: - data = line.rstrip() - if data.startswith("@"): - if not accs: - break - parse = False - header = data[1:].split(" ") - seq_name = header[0] - if seq_name in accs: - fq_dict[seq_name] = { - "sequence": "", - "description": " ".join(header[1:]), - "score": "", - } - parse = True - append_seq = "sequence" - accs.remove(seq_name) - elif parse: - if data == "+": - append_seq = "score" - else: - fq_dict[seq_name][append_seq] += data - else: - with open(fq_path, "r") as raw: - for line in raw: - data = line.rstrip() - if data.startswith("@"): - if not accs: - break - parse = False - header = data[1:].split(" ") - seq_name = header[0] - if seq_name in accs: - fq_dict[seq_name] = { - "sequence": "", - "description": " ".join(header[1:]), - "score": "", - } - parse = True - append_seq = "sequence" - accs.remove(seq_name) - elif parse: - if data == "+": - append_seq = "score" - else: - fq_dict[seq_name][append_seq] += data - return fq_dict - - -def cli(): - - parser = argparse.ArgumentParser(description="Inputs accession, extracts fastq") - parser.add_argument( - "-a", - "--accession", - help='"-" for stdin. For coordinates ' - + "append [$START-$END] - reverse coordinates for antisense", - ) - parser.add_argument("-i", "--input", help="File with accessions") - parser.add_argument("-f", "--fastq", help="FASTQ input", required=True) - args = parser.parse_args() - - if args.input: # input file - input_file = format_path(args.input) - with open(input_file, "r") as raw: - accs = [x.rstrip().split("\t")[args.column - 1] for x in raw if x.rstrip()] - else: # assume we are using accessions - if "-" == args.accession: # stdin - data = stdin2str() - accs = data.split() - else: - if {'"', "'"}.intersection(set(args.accession)): - args.accession = args.accession.replace('"', "").replace("'", "") - if "," in args.accession: - accs = args.accession.split(",") - elif re.search(r"\s", args.accession): - accs = args.accession.split() - else: - accs = [args.accession] - - fq_path = format_path(args.fastq) - fq_dict = acc2fq(fq_path, set(accs)) - fastq_str = dict2fq(fq_dict) - - print(fastq_str.rstrip(), flush=True) - sys.exit(0) - - -if __name__ == "__main__": - cli() diff --git a/mycotools/cli.py b/mycotools/cli.py new file mode 100644 index 0000000..8092307 --- /dev/null +++ b/mycotools/cli.py @@ -0,0 +1,62 @@ +#! /usr/bin/env python3 +"""Top-level dispatcher for the `mycotools` command. + +`mycotools` is the downstream-analysis entrypoint; the database lifecycle +(update/extract/predb/manage/accession/files) lives under the separate `mtdb` +command. Routes `mycotools ...` to a group dispatcher, or, for a leaf +tool like `rename`, directly to that tool's module.""" +from mycotools.lib.subcmd import Dispatcher + +# group name/alias -> submodule within the mycotools package. Groups (download, +# homology, cluster, phylo, stats, seq, gff) are subpackage dispatchers; +# `rename` is a single leaf module. +SUBCOMMANDS = { + "download": "download", + "d": "download", + "homology": "homology", + "h": "homology", + "cluster": "cluster", + "c": "cluster", + "phylo": "phylo", + "p": "phylo", + "rename": "rename", + "r": "rename", + "stats": "stats", + "seq": "seq", + "gff": "gff", +} + +DESCRIPTION = """Mycotools downstream-analysis toolkit + +Groups (all following arguments are forwarded to the group/tool): + download (d) download genomes/annotations from JGI or NCBI + homology (h) search query sequence(s) against the database or a fasta + cluster (c) cluster sequences / circumscribe homology groups + phylo (p) phylogenies and phylogenetic pipelines (crap/tree/synteny) + rename (r) substitute MTDB ome codes with taxonomic names in a file + stats annotation / assembly statistics + seq sequence & coordinate transforms + gff manipulate / render gff3 files + +Database operations (build/manage/query the MTDB) live under the `mtdb` command; +run `mtdb -h`. + +Examples: + mycotools download jgi -h + mycotools phylo crap -h + mycotools homology -h""" + +_dispatcher = Dispatcher( + "mycotools", + "mycotools", + SUBCOMMANDS, + DESCRIPTION, + metavar="GROUP", + arg_help="analysis group or tool (see below)", +) +main = _dispatcher.main +cli = _dispatcher.cli + + +if __name__ == "__main__": + cli() diff --git a/mycotools/cluster/__init__.py b/mycotools/cluster/__init__.py new file mode 100644 index 0000000..85987b7 --- /dev/null +++ b/mycotools/cluster/__init__.py @@ -0,0 +1,39 @@ +#! /usr/bin/env python3 +"""Dispatcher for the `mycotools cluster` subcommand. + +Routes `mycotools cluster ...` to a clustering module: `db` +circumscribes database sequences into homology groups, `fasta` runs iterative +sequence-similarity clustering of a fasta.""" +from mycotools.lib.subcmd import Dispatcher + +# subcommand name/alias -> submodule within this package (mycotools.cluster.) +SUBCOMMANDS = { + "db": "db", + "fasta": "fasta", + "fa": "fasta", +} + +DESCRIPTION = """Group sequences into clusters / homology groups + +Methods (all following arguments are forwarded to the method): + db circumscribe database sequences into homology groups + fasta (fa) iterative sequence-similarity clustering of a fasta + +Examples: + mycotools cluster db -h + mycotools cluster fasta -h""" + +_dispatcher = Dispatcher( + "mycotools cluster", + "mycotools.cluster", + SUBCOMMANDS, + DESCRIPTION, + metavar="METHOD", + arg_help="clustering method (see below)", +) +main = _dispatcher.main +cli = _dispatcher.cli + + +if __name__ == "__main__": + cli() diff --git a/mycotools/cluster/__main__.py b/mycotools/cluster/__main__.py new file mode 100644 index 0000000..66248d1 --- /dev/null +++ b/mycotools/cluster/__main__.py @@ -0,0 +1,6 @@ +#! /usr/bin/env python3 +"""Enable ``python -m mycotools.cluster`` to run the cluster dispatcher.""" +from mycotools.cluster import cli + +if __name__ == "__main__": + cli() diff --git a/mycotools/db2hgs.py b/mycotools/cluster/db.py similarity index 85% rename from mycotools/db2hgs.py rename to mycotools/cluster/db.py index e5f9eb5..39db58a 100755 --- a/mycotools/db2hgs.py +++ b/mycotools/cluster/db.py @@ -1,18 +1,21 @@ #! /usr/bin/env python3 -import os import sys import shutil +import logging import argparse import subprocess import multiprocessing as mp from statistics import stdev, StatisticsError from collections import defaultdict, Counter -from mycotools.acc2fa import dbmain as acc2fa -from mycotools.db2files import soft_main as symlink_files -from mycotools.lib.dbtools import mtdb, primaryDB -from mycotools.lib.biotools import fa2dict, dict2fa, fa2dict_accs -from mycotools.lib.kontools import format_path, eprint, mkOutput, findExecs +from mycotools.mtdb.files import soft_main as symlink_files +from mycotools.lib.dbtools import mtdb, primary_db +from mycotools.lib.biotools import dict2fa, fa2dict_accs +from mycotools.lib.kontools import format_path, mk_output, find_execs, setup_logging +from pathlib import Path + + +logger = logging.getLogger(__name__) def mk_db2hg_output(out_dir, nscg=False): @@ -22,14 +25,14 @@ def mk_db2hg_output(out_dir, nscg=False): nscg_dir = out_dir + "near_single_copy_genes/" hg_seq_dir = out_dir + "hgs/" - if not os.path.isdir(wrk_dir): - os.mkdir(wrk_dir) - if not os.path.isdir(scg_dir): - os.mkdir(scg_dir) - if not os.path.isdir(nscg_dir) and nscg: - os.mkdir(nscg_dir) - if not os.path.isdir(hg_seq_dir): - os.mkdir(hg_seq_dir) + if not Path(wrk_dir).is_dir(): + Path(wrk_dir).mkdir() + if not Path(scg_dir).is_dir(): + Path(scg_dir).mkdir() + if not Path(nscg_dir).is_dir() and nscg: + Path(nscg_dir).mkdir() + if not Path(hg_seq_dir).is_dir(): + Path(hg_seq_dir).mkdir() return wrk_dir, scg_dir, nscg_dir, hg_seq_dir @@ -46,12 +49,12 @@ def run_mmseqs( """Run MMseqs clustering by sym linking MTDB proteomes""" symlink_files(["faa"], db, wrk_dir, verbose=False) # symlink proteomes cluster_res_file = wrk_dir + "raw_hgs.tsv" - if not os.path.isfile(cluster_res_file): # NEED to add to log removal + if not Path(cluster_res_file).is_file(): # NEED to add to log removal # be cautious about shell injection because we need to glob int(cpus) float(min_id) float(min_cov) - if not os.path.isdir(wrk_dir): + if not Path(wrk_dir).is_dir(): raise OSError("invalid working directory") elif not algorithm in {"mmseqs easy-linclust", "mmseqs easy-cluster"}: raise OSError("invalid mmseqs binary") @@ -82,18 +85,18 @@ def run_mmseqs( ) # stderr = subprocess.DEVNULL) shutil.move(wrk_dir + "cluster_cluster.tsv", cluster_res_file) - elif os.path.getsize(cluster_res_file): + elif Path(cluster_res_file).stat().st_size: mmseqs_cmd = 0 else: mmseqs_cmd = 1 if mmseqs_cmd: - eprint("\tERROR: cluster failed") + logger.error("cluster failed") sys.exit(1) - if os.path.isfile(wrk_dir + "cluster_all_seqs.fasta"): - os.remove(wrk_dir + "cluster_all_seqs.fasta") - if os.path.isfile(wrk_dir + "cluster_rep_seq.fasta"): - os.remove(wrk_dir + "cluster_rep_seq.fasta") - if os.path.isdir(wrk_dir + "tmp/"): + if Path(wrk_dir + "cluster_all_seqs.fasta").is_file(): + Path(wrk_dir + "cluster_all_seqs.fasta").unlink() + if Path(wrk_dir + "cluster_rep_seq.fasta").is_file(): + Path(wrk_dir + "cluster_rep_seq.fasta").unlink() + if Path(wrk_dir + "tmp/").is_dir(): shutil.rmtree(wrk_dir + "tmp/") return cluster_res_file @@ -117,7 +120,6 @@ def parse_1to1(hg_file, useableOmes=set()): for hg, genes in hg2gene.items(): genes = [x for x in genes if x[: x.find("_")] in useableOmes] - omes = set([x[: x.find("_")] for x in genes]) hg2gene = { k: v for k, v in sorted(hg2gene.items(), key=lambda x: len(x[1]), reverse=True) } @@ -239,7 +241,7 @@ def write_hgs(hg, genes, wrk_dir, write_dir): with open(f"{write_dir}{hg}.faa.tmp", "w") as out: out.write(dict2fa(fa_dict)) - os.rename(f"{write_dir}{hg}.faa.tmp", f"{write_dir}{hg}.faa") + Path(f"{write_dir}{hg}.faa.tmp").rename(f"{write_dir}{hg}.faa") def align_hg(hg_fa, out_fa, cpus=1): @@ -248,7 +250,7 @@ def align_hg(hg_fa, out_fa, cpus=1): cmd = subprocess.call( ["mafft", "--auto", "--thread", f"-{cpus}", hg_fa], stdout=out ) # , stderr = subprocess.PIPE) - os.rename(out_fa + ".tmp", out_fa) + Path(out_fa + ".tmp").rename(out_fa) return cmd @@ -259,7 +261,7 @@ def hmmbuild_hg(msa_fa, out_hmm, cpus=1): stdout=subprocess.PIPE, stderr=subprocess.PIPE, ) - os.rename(out_hmm + ".tmp", out_hmm) + Path(out_hmm + ".tmp").rename(out_hmm) return cmd @@ -350,16 +352,16 @@ def main( aln_file = out_dir + "accessory_alignment.phy" wrk_dir, scg_dir, nscg_dir, hg_seq_dir = mk_db2hg_output(out_dir, nscg) - print("\nClustering protein sequences", flush=True) + logger.info("Clustering protein sequences") raw_hg_file = run_mmseqs(db, wrk_dir, algorithm, min_id, min_cov, sensitivity, cpus) useable_omes = set(db.keys()) - print("\nCompiling homology groups (HGs)", flush=True) + logger.info("Compiling homology groups (HGs)") ome_num, gene2hg, i2ome, hg2gene = compile_homolog_groups( raw_hg_file, hg_output, useable_omes ) - print("\nIdentifying single-copy HGs", flush=True) + logger.info("Identifying single-copy HGs") schgs, nschgs, hg2stats, full_hgs, hg2d_omes = id_near_schgs( hg2gene, set(i2ome), @@ -370,8 +372,8 @@ def main( min_genomes=min_genomes, ) - print("\nWriting output", flush=True) - ome2pan = pangenome_output(pan_file, aln_file, hg2gene, hg2d_omes, max_mis_ome=0) + logger.info("Writing output") + pangenome_output(pan_file, aln_file, hg2gene, hg2d_omes, max_mis_ome=0) with open(hg2missing_genome_file, "w") as out: out.write("#hg\tmissing\n") @@ -385,28 +387,28 @@ def main( {k: hg2stats[k] for k in list(full_hgs)}, full_hg_stats_file, sort=True ) if nscg: - print(f"\t{len(nschgs)} near single-copy HGs", flush=True) + logger.info(f"{len(nschgs)} near single-copy HGs") with mp.Pool(processes=cpus) as pool: pool.starmap( write_hgs, ( (hg, hg2gene[hg], wrk_dir, nscg_dir) for hg in nschgs - if not os.path.isfile(f"{nscg_dir}{hg}.faa") + if not Path(f"{nscg_dir}{hg}.faa").is_file() ), ) # for hg in nschgs: # if not os.path.isfile(f'{nscg_dir}{hg}.faa'): # write_hgs(db, hg, hg2gene[hg], wrk_dir, nscg_dir) if schgs: - print(f"\t{len(schgs)} single-copy HGs", flush=True) + logger.info(f"{len(schgs)} single-copy HGs") with mp.Pool(processes=cpus) as pool: pool.starmap( write_hgs, ( (hg, hg2gene[hg], wrk_dir, scg_dir) for hg in schgs - if not os.path.isfile(f"{scg_dir}{hg}.faa") + if not Path(f"{scg_dir}{hg}.faa").is_file() ), ) # for hg in schgs: @@ -419,7 +421,7 @@ def main( ( (hg, genes, wrk_dir, hg_seq_dir) for hg, genes in hg2genes.items() - if not os.path.isfile(f"{hg_seq_dir}{hg}.faa") + if not Path(f"{hg_seq_dir}{hg}.faa").is_file() ), ) @@ -428,25 +430,23 @@ def main( # write_hgs(db, hg, genes, wrk_dir, hg_seq_dir) if hmm: - print("\nAligning and building HMMs", flush=True) + logger.info("Aligning and building HMMs") msa_dir = out_dir + "msa/" hmm_dir = out_dir + "hmm/" for d in [msa_dir, hmm_dir]: - if not os.path.isdir(d): - os.mkdir(d) + if not Path(d).is_dir(): + Path(d).mkdir() if nscg: srch_hgs = nschgs - hg_dir = nscg_dir else: srch_hgs = schgs - hg_dir = scg_dir for hg in srch_hgs: - if not os.path.isfile(f"{msa_dir}{hg}.mafft.faa"): - mafft_code = align_hg( + if not Path(f"{msa_dir}{hg}.mafft.faa").is_file(): + align_hg( f"{nscg_dir}{hg}.faa", f"{msa_dir}{hg}.mafft.faa", cpus=cpus ) - if not os.path.isfile(f"{hmm_dir}{hg}.hmm"): - hmm_code = hmmbuild_hg( + if not Path(f"{hmm_dir}{hg}.hmm").is_file(): + hmmbuild_hg( f"{msa_dir}{hg}.mafft.faa", f"{hmm_dir}{hg}.hmm", cpus=cpus ) @@ -455,7 +455,7 @@ def cli(): parser = argparse.ArgumentParser( description="Circumscribe protein sequences into homology groups" ) - parser.add_argument("-d", "--mtdb", default=primaryDB()) + parser.add_argument("-d", "--mtdb", default=primary_db()) parser.add_argument( "-m", "--mean", @@ -492,25 +492,26 @@ def cli(): parser.add_argument("-o", "--out_dir", help="WARNING: will not overwrite") parser.add_argument("-c", "--cpus", type=int, default=mp.cpu_count()) args = parser.parse_args() + setup_logging(verbose=getattr(args, "verbose", False)) if args.min_genomes > 1 or args.min_genomes < 0: - eprint("\nERROR: --min_genomes must be between 0 and 1", flush=True) + logger.error("--min_genomes must be between 0 and 1") sys.exit(143) db = mtdb(format_path(args.mtdb)) if args.out_dir: out_dir = format_path(args.out_dir) - if not os.path.isdir(out_dir): - os.mkdir(out_dir) + if not Path(out_dir).is_dir(): + Path(out_dir).mkdir() out_dir += "/" else: - out_dir = mkOutput(args.out_dir, "db2hgs") + out_dir = mk_output(args.out_dir, "cluster_db") execs = ["mmseqs"] if args.hmm: execs.extend(["hmmbuild", "mafft"]) - findExecs(execs, set(execs)) + find_execs(execs, set(execs)) main( db, diff --git a/mycotools/fa2clus.py b/mycotools/cluster/fasta.py similarity index 82% rename from mycotools/fa2clus.py rename to mycotools/cluster/fasta.py index c1df052..2d8f8f2 100755 --- a/mycotools/fa2clus.py +++ b/mycotools/cluster/fasta.py @@ -5,31 +5,28 @@ # NEED to try MCL using the binary # NEED to make rerunning aggclus not use old data -import os +import logging import re import sys import copy import shutil -import string -import random -import tempfile import argparse -import itertools import subprocess import pandas as pd from collections import defaultdict from mycotools.lib.kontools import ( - multisub, - findExecs, + find_execs, format_path, - eprint, - vprint, read_json, write_json, - mkOutput, + mk_output, fmt_float, + setup_logging, ) from mycotools.lib.biotools import fa2dict, dict2fa +from pathlib import Path + +logger = logging.getLogger(__name__) sys.setrecursionlimit(1000000) @@ -53,7 +50,7 @@ def run_mmseqs( verbose=False, ): res_path = res_base + "_cluster.tsv" - tmp_dir = os.path.dirname(res_path) + "/tmp/" + tmp_dir = str(Path(res_path).parent) + "/tmp/" if verbose: stdout, stderr = None, None else: @@ -91,11 +88,11 @@ def run_mmseqs( raise ClusteringError( "Clustering failed: " + str(mmseqs_exit) + " " + " ".join(mmseqs_cmd) ) - if os.path.isfile(res_base + "_all_seqs.fasta"): - os.remove(res_base + "_all_seqs.fasta") - if os.path.isfile(res_base + "_rep_seq.fasta"): - os.remove(res_base + "_rep_seq.fasta") - if os.path.isdir(tmp_dir): + if Path(res_base + "_all_seqs.fasta").is_file(): + Path(res_base + "_all_seqs.fasta").unlink() + if Path(res_base + "_rep_seq.fasta").is_file(): + Path(res_base + "_rep_seq.fasta").unlink() + if Path(tmp_dir).is_dir(): shutil.rmtree(tmp_dir) return res_path @@ -122,10 +119,10 @@ def parse_mmseqs_clus(res_path): return hg2gene, gene2hg -def makeDmndDB(diamond, queryFile, output_dir, cpus=1): +def make_dmnd_db(diamond, queryFile, output_dir, cpus=1): """create a diamond database: diamond: binary path, queryFile: query_path""" - outputDB = output_dir + re.sub(r"\.[^\.]+$", "", os.path.basename(queryFile)) + outputDB = output_dir + re.sub(r"\.[^\.]+$", "", Path(queryFile).name) cmd = [ diamond, "makedb", @@ -142,7 +139,7 @@ def makeDmndDB(diamond, queryFile, output_dir, cpus=1): return outputDB, dmndDBcode -def runDmnd( +def run_dmnd( diamond, queryFile, queryDB, @@ -243,68 +240,6 @@ def rd_dmnd_distmtx(outputFile, minVal, pid=True): return distance_matrix # , outMatrix -def runUsearch(fasta, output, clus_var, cpus=1, verbose=False): - - if verbose: - subprocess.call( - [ - "usearch", - "-calc_distmx", - fasta, - "-tabbedout", - output, - "-clus_var", - clus_var, - "-threads", - str(cpus), - ] # stdout = subprocess.PIPE, - # stderr = subprocess.PIPE - ) - else: - subprocess.call( - [ - "usearch", - "-calc_distmx", - fasta, - "-tabbedout", - output, - "-clus_var", - clus_var, - "-threads", - str(cpus), - ], - stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL, - ) - - -def rd_usrch_distmtx(dis_path, sep="\t"): - """Imports a distance matrix with each line formatted as `organism $SEP organism $SEP distance`. - - this is equivalent to the `-tabbedout` argument in `usearch -calc_distmx`. The function - compiles a dictionary with the information and reciprocal information for each organism in each - line, then converts this dictionary of dictionaries into a pandas dataframe. As a distance matrix, - NA values are converted to maximum distance (1).""" - - distance_matrix = pd.DataFrame() - with open(dis_path, "r") as raw: - data = raw.read() - prepData = [x.rstrip() for x in data.split("\n") if x != ""] - - dist_dict = {} - for line in prepData: - vals = line.split(sep) - if vals[0] not in dist_dict: - dist_dict[vals[0]] = {} - if vals[1] not in dist_dict: - dist_dict[vals[1]] = {} - dist_dict[vals[0]][vals[1]] = float(vals[2]) - dist_dict[vals[1]][vals[0]] = float(vals[2]) - - distance_matrix = pd.DataFrame(dist_dict).sort_index(0).sort_index(1) - - return distance_matrix.fillna(1.0) - - def scikitaggd(distance_matrix, maxDist=0.6, linkage="single"): """Performs agglomerative clustering and extracts the cluster labels, then sorts according to cluster number. In the future, this will also extract a Newick tree.""" @@ -323,7 +258,7 @@ def scikitaggd(distance_matrix, maxDist=0.6, linkage="single"): return clusters -def getNewick(node, newick, parentdist, leaf_names): +def get_newick(node, newick, parentdist, leaf_names): """Code from https://stackoverflow.com/questions/28222179/save-dendrogram-to-newick-format adopted from @jfn""" if node.is_leaf(): @@ -333,13 +268,13 @@ def getNewick(node, newick, parentdist, leaf_names): newick = "):%.2f%s" % (parentdist - node.dist, newick) else: newick = ");" - newick = getNewick(node.get_left(), newick, node.dist, leaf_names) - newick = getNewick(node.get_right(), ",%s" % (newick), node.dist, leaf_names) + newick = get_newick(node.get_left(), newick, node.dist, leaf_names) + newick = get_newick(node.get_right(), ",%s" % (newick), node.dist, leaf_names) newick = "(%s" % (newick) return newick -def getClusterLabels(labels, clusters): +def get_cluster_labels(labels, clusters): i = iter(labels) protoclusters = {next(i): x for x in clusters} @@ -350,7 +285,7 @@ def getClusterLabels(labels, clusters): return clusters -def runMCL(distMat, inflation): +def run_mcl(distMat, inflation): key2gene = {i: v for i, v in enumerate(list(distMat.keys()))} npMat = distMat.to_numpy() @@ -372,14 +307,14 @@ def scipyaggd(distMat, maxDist, method="single"): linkage_matrix = hierarchy.linkage(squareform_matrix, method) tree = hierarchy.to_tree(linkage_matrix) fcluster = hierarchy.fcluster(linkage_matrix, maxDist, criterion="distance") - clusters = getClusterLabels(distMat.index, fcluster) + clusters = get_cluster_labels(distMat.index, fcluster) return clusters, tree def dmnd_main(fa_path, minVal, output_dir, distFile, pid=True, verbose=False, cpus=1): - queryDB, makeDBcode = makeDmndDB("diamond", fa_path, output_dir, cpus=cpus) - dmndOut, dmndCode = runDmnd( + queryDB, makeDBcode = make_dmnd_db("diamond", fa_path, output_dir, cpus=cpus) + dmndOut, dmndCode = run_dmnd( "diamond", fa_path, queryDB, @@ -395,13 +330,13 @@ def dmnd_main(fa_path, minVal, output_dir, distFile, pid=True, verbose=False, cp def usrch_main(fasta, min_id, output, cpus=1, verbose=False): - vprint("\nusearch aligning", flush=True, v=verbose) + logger.debug("usearch aligning") runUsearch(fasta, output + ".dist", str(1 - min_id), cpus, verbose) distance_matrix = rd_usrch_distmtx(output + ".dist") return distance_matrix -def readLog(log_path, newLog): +def read_log(log_path, newLog): oldLog = read_json(log_path) if ( oldLog["fasta"] == newLog["fasta"] @@ -415,8 +350,8 @@ def readLog(log_path, newLog): for i in newLog["successes"]: i["cluster"] = tuple(i["cluster"]) elif newLog["distance_matrix"]: - if os.path.isfile(newLog["distance_matrix"]): - os.remove(newLog["distance_matrix"]) + if Path(newLog["distance_matrix"]).is_file(): + Path(newLog["distance_matrix"]).unlink() return newLog @@ -470,7 +405,7 @@ def cluster_iter_mmseqs( ) res_info = res_base + name res_path = res_info + "_cluster.tsv" - if not os.path.isfile(res_path): + if not Path(res_path).is_file(): run_mmseqs( params["fa"], res_info, @@ -485,17 +420,15 @@ def cluster_iter_mmseqs( cluster = cluster_dict[clusters[focal_gene]] focal_len = len(cluster) - vprint( - "\nITERATION " + logger.debug( + "ITERATION " + str(attempt) + ": " + focal_gene + " cluster size: " - + str(focal_len), - flush=True, - v=verbose, + + str(focal_len) ) - vprint("Cluster parameter: " + str(clus_var), flush=True, v=verbose) + logger.debug("Cluster parameter: " + str(clus_var)) iteration_dict = { "size": focal_len, "cluster_variable": clus_var, @@ -506,7 +439,7 @@ def cluster_iter_mmseqs( if focal_len >= min_seq: # if greater than minimum sequences if max_seq: # if there is a max set of sequences if focal_len <= max_seq: # if less than max sequences - vprint(spacer + "\tSUCCESS!", flush=True, v=verbose) + logger.debug(spacer + "SUCCESS!") exit_code = 0 direction = -1 log_dict["successes"].append(iteration_dict) @@ -514,7 +447,7 @@ def cluster_iter_mmseqs( else: # descend direction = 1 else: # no max sequences, minimum is met, this was successful - vprint(spacer + "\tSUCCESS!", flush=True, v=verbose) + logger.debug(spacer + "SUCCESS!") exit_code = 0 log_dict["successes"].append(iteration_dict) log_dict["successes"] = sort_iterations(log_dict["successes"]) @@ -536,11 +469,10 @@ def cluster_iter_mmseqs( exit_code = 0 break else: - eprint( + logger.warning( spacer - + "WARNING: Overshot - " - + "could not find parameters using current interval", - flush=True, + + "Overshot - " + + "could not find parameters using current interval" ) iteration = extract_closest_cluster( log_dict["iterations"], min_seq, max_seq @@ -553,7 +485,6 @@ def cluster_iter_mmseqs( exit_code = 1 # clus_const += (interval*direction) - ofocal_len = copy.copy(focal_len) oclus_var = copy.copy(clus_var) clus_var += interval * direction @@ -608,12 +539,12 @@ def cluster_iter_aggclus( clusters, tree = scipyaggd( params["dist"], float(clus_var), params["link"] ) # cluster - newick = getNewick(tree, "", tree.dist, list(params["dist"].index)) + newick = get_newick(tree, "", tree.dist, list(params["dist"].index)) write_data(newick, clusters, res_base + name) cluster_dict = defaultdict( list - ) # could be more efficient by grabbing in getClusterLabels + ) # could be more efficient by grabbing in get_cluster_labels for gene, index in clusters.items(): # create a dictionary clusID: # [genes] cluster_dict[index].append(gene) @@ -626,17 +557,15 @@ def cluster_iter_aggclus( cluster = None newick = "" - vprint( - "\nITERATION " + logger.debug( + "ITERATION " + str(attempt) + ": " + focal_gene + " cluster size: " - + str(focal_len), - flush=True, - v=verbose, + + str(focal_len) ) - vprint("Cluster parameter: " + str(clus_var), flush=True, v=verbose) + logger.debug("Cluster parameter: " + str(clus_var)) iteration_dict = { "size": focal_len, "cluster_variable": clus_var, @@ -648,7 +577,7 @@ def cluster_iter_aggclus( if focal_len >= min_seq: # if greater than minimum sequences if max_seq: # if there is a max set of sequences if focal_len <= max_seq: # if less than max sequences - vprint(spacer + "\tSUCCESS!", flush=True, v=verbose) + logger.debug(spacer + "SUCCESS!") exit_code = 0 direction = -1 log_dict["successes"].append(iteration_dict) @@ -656,7 +585,7 @@ def cluster_iter_aggclus( else: # descend direction = 1 else: # no max sequences, minimum is met, this was successful - vprint(spacer + "\tSUCCESS!", flush=True, v=verbose) + logger.debug(spacer + "SUCCESS!") exit_code = 0 log_dict["successes"].append(iteration_dict) log_dict["successes"] = sort_iterations(log_dict["successes"]) @@ -673,11 +602,10 @@ def cluster_iter_aggclus( newick = log_dict["successes"][0]["tree"] exit_code = 0 else: - eprint( + logger.warning( spacer - + "WARNING: Overshot - " - + "could not find parameters using current interval", - flush=True, + + "Overshot - " + + "could not find parameters using current interval" ) iteration = extract_closest_cluster( log_dict["iterations"], min_seq, max_seq @@ -730,8 +658,8 @@ def main( max_var=1, ): - if not os.path.isdir(os.path.dirname(output) + "/working/"): - os.mkdir(os.path.dirname(output) + "/working/") + if not Path(str(Path(output).parent) + "/working/").is_dir(): + Path(str(Path(output).parent) + "/working/").mkdir() if search_program in {"usearch", "diamond"}: algorithm = "hierarchical" @@ -756,8 +684,8 @@ def main( "successes": [], } if log_path: - if os.path.isfile(log_path): - log_dict = readLog(log_path, log_dict) + if Path(log_path).is_file(): + log_dict = read_log(log_path, log_dict) write_json(log_dict, log_path) if algorithm == "hierarchical": @@ -769,9 +697,9 @@ def main( param_dict = { "dist": None, "link": linkage, - "dir": os.path.dirname(output) + "/", + "dir": str(Path(output).parent) + "/", } - if os.path.isfile(log_dict["distance_matrix"]): + if Path(log_dict["distance_matrix"]).is_file(): if search_program == "usearch": param_dict["dist"] = rd_usrch_distmtx(log_dict["distance_matrix"]) else: @@ -781,7 +709,7 @@ def main( else: if search_program == "diamond": # elif if above lines not highlighted if not dmnd_dir: - dmnd_dir = os.path.dirname(log_dict["distance_matrix"]) + "/" + dmnd_dir = str(Path(log_dict["distance_matrix"]).parent) + "/" param_dict["dist"] = dmnd_main( fa_path, clus_const, @@ -800,11 +728,11 @@ def main( "fa": fa_path, "bin": search_program, "clus_const": clus_const, - "dir": os.path.dirname(output) + "/", + "dir": str(Path(output).parent) + "/", } if focal_gene: - vprint("\nClustering", flush=True, v=verbose) + logger.debug("Clustering") if algorithm == "hierarchical": cluster, newick, log_dict, error = cluster_iter_aggclus( param_dict, @@ -853,14 +781,12 @@ def main( if focal_gene: res_base = param_dict["dir"] + focal_gene else: - res_base = param_dict["dir"] + re.sub( - r"\.[^\.]+$", "", os.path.basename(fa_path) - ) + res_base = param_dict["dir"] + re.sub(r"\.[^\.]+$", "", Path(fa_path).name) if algorithm == "hierarchical": clusters, tree = scipyaggd( param_dict["dist"], float(clus_var), param_dict["link"] ) # cluster - newick = getNewick(tree, "", tree.dist, list(param_dict["dist"].index)) + newick = get_newick(tree, "", tree.dist, list(param_dict["dist"].index)) write_data(newick, clusters, res_base) return clusters, None, None, log_dict else: @@ -970,6 +896,7 @@ def cli(): parser.add_argument("-c", "--cpus", default=1, type=int) parser.add_argument("-v", "--verbose", action="store_true") args = parser.parse_args() + setup_logging(verbose=getattr(args, "verbose", False)) # if args.refine and not args.max_seq: # eprint('\nERROR: --max_seq required for refinement', flush = True) @@ -981,10 +908,10 @@ def cli(): "centroid", "single", }: - eprint("\nERROR: Invalid linkage criterium", flush=True) + logger.error("Invalid linkage criterium") sys.exit(1) elif args.distance_type not in {"identity", "bitscore"}: - eprint("\nERROR: Invalid distance type", flush=True) + logger.error("Invalid distance type") sys.exit(3) if args.distance_type == "identity": pid = True @@ -1007,13 +934,13 @@ def cli(): if not args.cluster_variable: args.cluster_variable = 0.3 else: - eprint("\nERROR: Invalid alignment software", flush=True) + logger.error("Invalid alignment software") sys.exit(2) - findExecs([args.alignment.split()[0]], exit=set(args.alignment.split()[0])) + find_execs([args.alignment.split()[0]], exit=set(args.alignment.split()[0])) interval = args.interval if args.interval < 0 or args.interval > 1: - eprint("\nERROR: --interval must be between 0 and 1", flush=True) + logger.error("--interval must be between 0 and 1") # if args.iterative and args.alignment in {'diamond', 'usearch'}: # clus_var = 1 - args.cluster_constant @@ -1021,9 +948,8 @@ def cli(): 1 - args.cluster_variable <= args.cluster_constant and args.alignment != "mmseqs" ): - eprint( - "\nWARNING: 1 - maximum distance exceeds minimum connection, clustering is ineffective", - flush=True, + logger.warning( + "1 - maximum distance exceeds minimum connection, clustering is ineffective" ) sys.exit(3) else: @@ -1035,20 +961,20 @@ def cli(): fa = fa2dict(fa_path) if args.iterative: if len(fa) < args.min_seq: - eprint("\nERROR: minimum sequences is greater than fasta input", flush=True) + logger.error("minimum sequences is greater than fasta input") sys.exit(5) elif args.iterative not in fa: - eprint("\nERROR: " + args.iterative + " not in " + fa_path, flush=True) + logger.error("" + args.iterative + " not in " + fa_path) sys.exit(6) if args.output: - if not os.path.isdir(format_path(args.output)): - os.mkdir(format_path(args.output)) + if not Path(format_path(args.output)).is_dir(): + Path(format_path(args.output)).mkdir() dmnd_dir = format_path(args.output) - output = dmnd_dir + re.sub(r"\.[^\.]+$", "", os.path.basename(fa_path)) + output = dmnd_dir + re.sub(r"\.[^\.]+$", "", Path(fa_path).name) else: - dmnd_dir = mkOutput(os.getcwd() + "/", "fa2clus") - output = dmnd_dir + re.sub(r"\.[^\.]+$", "", os.path.basename(fa_path)) + dmnd_dir = mk_output(str(Path.cwd()) + "/", "cluster_fasta") + output = dmnd_dir + re.sub(r"\.[^\.]+$", "", Path(fa_path).name) cluster, tree, overshot, log_dict = main( fa_path, @@ -1062,7 +988,7 @@ def cli(): focal_gene=args.iterative, interval=interval, output=output, - log_path=dmnd_dir + "." + os.path.basename(fa_path) + ".fa2clus.json", + log_path=dmnd_dir + "." + Path(fa_path).name + ".fa2clus.json", pid=pid, dmnd_dir=dmnd_dir, cpus=args.cpus, @@ -1077,7 +1003,7 @@ def cli(): try: output_fa = {x: input_fa[x] for x in cluster} except TypeError: - eprint("\nERROR: empty cluster", flush=True) + logger.error("empty cluster") sys.exit(10) with open(output + ".fa", "w") as out: out.write(dict2fa(output_fa)) diff --git a/mycotools/deprecated.py b/mycotools/deprecated.py new file mode 100644 index 0000000..136857a --- /dev/null +++ b/mycotools/deprecated.py @@ -0,0 +1,77 @@ +#! /usr/bin/env python3 +"""Backwards-compatibility shims for the pre-v2 flat command names. + +Every tool that moved under the nested `mycotools`/`mtdb` dispatchers keeps its +old flat console command (e.g. `jgiDwnld`, `acc2fa`) working for one transition +period. Each old command resolves to a shim here that prints a deprecation +notice to stderr and then hands off to the tool's relocated `cli()` unchanged; +`sys.argv` is untouched, so argument parsing behaves exactly as before. + +These direct entrypoints are TEMPORARY and will be removed in a subsequent +release -- use the nested invocation printed in the warning instead.""" +import sys +import importlib + + +def _make_shim(old, new, target): + """Build a console-script shim: warn, then delegate to `target`'s cli(). + + old : the deprecated flat command name + new : the nested invocation the user should switch to + target : importable module exposing `cli()` (the relocated tool) + """ + + def shim(): + sys.stderr.write( + f"WARNING: `{old}` is a deprecated entrypoint and will be removed in " + f"a subsequent release. Use `{new}` instead.\n" + ) + importlib.import_module(target).cli() + + shim.__name__ = old + shim.__qualname__ = old + shim.__doc__ = f"Deprecated alias for `{new}`." + return shim + + +# --- relocated under the `mycotools` analysis entrypoint --------------------- +jgiDwnld = _make_shim("jgiDwnld", "mycotools download jgi", "mycotools.download.jgi") +ncbiDwnld = _make_shim( + "ncbiDwnld", "mycotools download ncbi", "mycotools.download.ncbi" +) +db2search = _make_shim("db2search", "mycotools homology db", "mycotools.homology.db") +fa2hmmer2fa = _make_shim( + "fa2hmmer2fa", "mycotools homology fasta", "mycotools.homology.fasta" +) +db2hgs = _make_shim("db2hgs", "mycotools cluster db", "mycotools.cluster.db") +fa2clus = _make_shim("fa2clus", "mycotools cluster fasta", "mycotools.cluster.fasta") +crap = _make_shim("crap", "mycotools phylo crap", "mycotools.phylo.crap") +fa2tree = _make_shim("fa2tree", "mycotools phylo tree", "mycotools.phylo.tree") +db2microsyntree = _make_shim( + "db2microsyntree", "mycotools phylo synteny", "mycotools.phylo.synteny" +) +annotationStats = _make_shim( + "annotationStats", "mycotools stats annotation", "mycotools.stats.annotation" +) +assemblyStats = _make_shim( + "assemblyStats", "mycotools stats assembly", "mycotools.stats.assembly" +) +fna2faa = _make_shim("fna2faa", "mycotools seq translate", "mycotools.seq.translate") +coords2fa = _make_shim("coords2fa", "mycotools seq coords", "mycotools.seq.coords") +gff2seq = _make_shim("gff2seq", "mycotools seq gff", "mycotools.seq.gff") +bioreform = _make_shim("bioreform", "mycotools seq convert", "mycotools.seq.convert") +fa2mass = _make_shim("fa2mass", "mycotools seq mass", "mycotools.seq.mass") +add2gff = _make_shim("add2gff", "mycotools gff add", "mycotools.gff.add") +gff2svg = _make_shim("gff2svg", "mycotools gff svg", "mycotools.gff.svg") + +# --- relocated under the `mtdb` database entrypoint -------------------------- +db2files = _make_shim("db2files", "mtdb files", "mycotools.mtdb.files") + +# --- relocated under the `mycotools` analysis entrypoint --------------------- +ome2name = _make_shim("ome2name", "mycotools rename", "mycotools.rename") + +# --- restored acc2* aliases (retired in v2, now `mtdb accession `) ---- +acc2fa = _make_shim("acc2fa", "mtdb accession fa", "mycotools.mtdb.acc2.fa") +acc2gff = _make_shim("acc2gff", "mtdb accession gff", "mycotools.mtdb.acc2.gff") +acc2gbk = _make_shim("acc2gbk", "mtdb accession gbk", "mycotools.mtdb.acc2.gbk") +acc2locus = _make_shim("acc2locus", "mtdb accession locus", "mycotools.mtdb.acc2.locus") diff --git a/mycotools/download/__init__.py b/mycotools/download/__init__.py new file mode 100644 index 0000000..a271b17 --- /dev/null +++ b/mycotools/download/__init__.py @@ -0,0 +1,38 @@ +#! /usr/bin/env python3 +"""Dispatcher for the `mycotools download` subcommand. + +Routes `mycotools download ...` to a per-source retrieval module. Note +that `mtdb update` already assimilates JGI/NCBI genomes into the primary MTDB; +these commands are for fetching genomes/annotations directly as files.""" +from mycotools.lib.subcmd import Dispatcher + +# subcommand name/alias -> submodule within this package (mycotools.download.) +SUBCOMMANDS = { + "jgi": "jgi", + "ncbi": "ncbi", +} + +DESCRIPTION = """Download genomes/annotations from external repositories + +Sources (all following arguments are forwarded to the source's downloader): + jgi download from JGI MycoCosm + ncbi download from NCBI GenBank/RefSeq + +Examples: + mycotools download jgi -h + mycotools download ncbi -h""" + +_dispatcher = Dispatcher( + "mycotools download", + "mycotools.download", + SUBCOMMANDS, + DESCRIPTION, + metavar="SOURCE", + arg_help="download source (see below)", +) +main = _dispatcher.main +cli = _dispatcher.cli + + +if __name__ == "__main__": + cli() diff --git a/mycotools/download/__main__.py b/mycotools/download/__main__.py new file mode 100644 index 0000000..c557d62 --- /dev/null +++ b/mycotools/download/__main__.py @@ -0,0 +1,6 @@ +#! /usr/bin/env python3 +"""Enable ``python -m mycotools.download`` to run the download dispatcher.""" +from mycotools.download import cli + +if __name__ == "__main__": + cli() diff --git a/mycotools/download/jgi.py b/mycotools/download/jgi.py new file mode 100755 index 0000000..cb278c4 --- /dev/null +++ b/mycotools/download/jgi.py @@ -0,0 +1,1054 @@ +#! /usr/bin/env python3 +""" +Download MycoCosm (JGI fungal) genome data. + +JGI retired its legacy ``get-directory`` XML download endpoint. Downloads now go +through the JGI Data Portal API (https://files.jgi.doe.gov), which this module +drives directly (no Globus): + + 1. authenticate at signon.jgi.doe.gov (the ``jgi_session`` cookie value is + the session token); + 2. list an organism's files via the ``mycocosm_file_list`` search endpoint; + 3. restore any archived (PURGED, on-tape) files via ``request_archived_files``; + 4. download immediately-available (RESTORED) files as a single zip stream via + the ``download_files`` endpoint, authorized with + ``Authorization: Bearer ``. + +The GFF3 hierarchy selects the *filtered* gene models only - JGI ``jat_label`` +``genes_filtered`` (i.e. the GeneCatalog / FilteredModels ``.gff``) - never the +unfiltered ``genes_all`` models. + +PLEASE respect JGI's rate limits. +""" + +import os +import re +import sys +import time +import shutil +import zipfile +import logging +import argparse +import requests +import pandas as pd +from tqdm import tqdm +from urllib.parse import unquote +from mycotools.lib.kontools import format_path, outro, intro, setup_logging +from mycotools.lib.dbtools import login_check +from pathlib import Path + +logger = logging.getLogger(__name__) + +# JGI Data Portal API endpoints (non-Globus) +SIGNON_URL = "https://signon.jgi.doe.gov/signon/create" +SEARCH_URL = "https://files.jgi.doe.gov/mycocosm_file_list/" +RESTORE_URL = "https://files.jgi.doe.gov/request_archived_files/" +DOWNLOAD_URL = "https://files-download.jgi.doe.gov/download_files/" + + +# =========================================================================== +# JGI Data Portal API - non-Globus download implementation (see module docstring) +# =========================================================================== + +# Ordered, most-preferred-first jat_labels and acceptable (un-gzipped) file +# formats per download type. Most importantly, "gff3" resolves only to the +# filtered gene models (genes_filtered), never genes_all. +_TYPE_LABELS = { + "gff3": (["genes_filtered"], {"gff", "gff3"}), + "faa": (["proteins_filtered"], {"fasta", "fa", "aa"}), + "transcript": (["transcripts_filtered"], {"fasta", "fa", "fna", "fsa", "nt"}), + "est": (["ests", "est_clusters"], {"fasta", "fa", "fna", "fsa"}), +} + +# Mitochondrial files that MycoCosm misfiles under a nuclear label. Mito +# assemblies are routinely tagged `assembly_unmasked` and shelved under "Genome +# Assembly (unmasked)" (e.g. Suilu4_MitoAssemblyScaffolds.fasta.gz), and mito +# annotations are tagged `genes_filtered` (e.g. Lst7536_1_MitoGenes.gff3.gz). +# Some portals list one and the same file under both the nuclear label and +# `assembly_mitochondrial`, so the label cannot discriminate and the filename +# has to. A mitochondrion is not an organismal genome and must never stand in +# for one - it is ~1/1000th the size, so the substitution silently produces a +# nonsense MTDB entry rather than an obvious failure. +# +# "mito" must open a filename token and be followed either by the rest of +# "mitochondri*" or by an assembly/annotation noun. Both guards protect nuclear +# files whose organism name merely contains the substring: it appears mid-token +# in Fomitopsis_*_AssemblyScaffolds.fasta.gz and token-initially in +# Mitosporidium_*_AssemblyScaffolds.fasta.gz. +_MITO_FILE = re.compile( + r"(?:^|[_.\-])mito" + r"(?:chondri\w*|(?=[_.\-]?(?:assembl|scaffold|contig|chromosom|genome|gene)))", + re.IGNORECASE, +) + + +def _file_format(f): + """Return a file's lowercase format from its metadata, falling back to the + filename extension (ignoring a trailing .gz).""" + fmt = ((f.get("metadata") or {}).get("file_format") or "").lower() + if fmt: + return fmt + name = re.sub(r"\.gz$", "", f.get("file_name", ""), flags=re.IGNORECASE) + ext = re.search(r"\.([^.]+)$", name) + return ext[1].lower() if ext else "" + + +def _jat_label(f): + """Return a file's lowercase JGI Analysis Task label (the canonical file + role, e.g. assembly_masked, genes_filtered).""" + return ((f.get("metadata") or {}).get("jat_label") or "").lower() + + +def _is_restored(f): + """Whether a file is immediately downloadable (on disk) rather than PURGED + to tape.""" + return str(f.get("file_status", "")).upper() == "RESTORED" + + +def _is_mito(f): + """Whether a file is a mitochondrial assembly/annotation rather than the + organismal one, judged by filename because JGI's labels misreport it (see + ``_MITO_FILE``).""" + return bool(_MITO_FILE.search(f.get("file_name", ""))) + + +def select_file(files, ftype, masked=True): + """Choose the single best file record for `ftype` from an organism's file + list. + + Mitochondrial files are excluded outright, whatever their label - see + ``_MITO_FILE``. Of the remainder, preference goes in order to: + 1. immediately-available (RESTORED) files over archived (PURGED) ones; + 2. the type's own label preference (e.g. masked assembly before unmasked + when ``masked`` is set). + + Returns the chosen file dict, or None if the organism has no matching file. + """ + if ftype == "fna": + labels = ( + ["assembly_masked", "assembly_unmasked"] + if masked + else ["assembly_unmasked", "assembly_masked"] + ) + formats = {"fasta", "fa", "fna", "fsa"} + name_ok = lambda n: True + elif ftype in _TYPE_LABELS: + labels, formats = _TYPE_LABELS[ftype] + if ftype == "faa": + # proteins_filtered also tags the .tab annotation and promoter files; + # keep only the actual proteome fasta + name_ok = lambda n: bool( + re.search(r"\.aa\.fa(sta)?(\.gz)?$", n, re.IGNORECASE) + ) + else: + name_ok = lambda n: True + else: + return None + + candidates = [] + for f in files: + label = _jat_label(f) + if label not in labels: + continue + if _is_mito(f): + continue + if _file_format(f) not in formats: + continue + if not name_ok(f.get("file_name", "")): + continue + status_rank = 0 if _is_restored(f) else 1 + candidates.append((status_rank, labels.index(label), f)) + if not candidates: + return None + candidates.sort(key=lambda t: (t[0], t[1])) + return candidates[0][2] + + +def parse_org_name(name): + """Split a JGI organism name (e.g. "Acaromyces ingoldii MCA 4198 v1.0") into + (genus, species, strain), dropping a trailing version token (strain is the + remaining words concatenated).""" + parts = str(name).split() + if not parts: + return "", "", "" + genus = parts[0] + species = parts[1] if len(parts) > 1 else "sp." + rest = parts[2:] + if rest and re.fullmatch(r"[vV]?\d+(\.\d+)*", rest[-1]): + rest = rest[:-1] + return genus, species, "".join(rest) + + +def _fill_if_empty(df, i, col, value): + """Set df.at[i, col] = value only when there is no existing non-empty value, + so curated reference genus/species/strain are preserved while bare-accession + input is populated from JGI metadata.""" + if not value: + return + cur = df.at[i, col] if col in df.columns else None + if ( + cur is None + or (isinstance(cur, float) and pd.isna(cur)) + or str(cur).strip() == "" + ): + df.at[i, col] = value + + +def jgi_api_login(user, pwd, max_attempts=5, spacer="\t"): + """Authenticate against JGI's signon service and return (session, token). + The session token (the jgi_session cookie value) authorizes the search, + restore, and download endpoints. Exits (100) after repeated failures.""" + session = requests.Session() + for attempt in range(1, max_attempts + 1): + try: + resp = session.post( + SIGNON_URL, data={"login": user, "password": pwd}, timeout=120 + ) + except requests.RequestException as error: + logger.warning(f"{spacer}\tJGI login error (attempt {attempt}): {error}") + time.sleep(5) + continue + token = session.cookies.get("jgi_session") + if resp.status_code == 200 and token: + return session, unquote(token) + logger.warning( + f"{spacer}\tJGI login failed (attempt {attempt}, status {resp.status_code})" + ) + time.sleep(5) + logger.error(f"{spacer}Failed {max_attempts} JGI login attempts.") + sys.exit(100) + + +def search_organism(session, portal_id, spacer="\t", max_attempts=3): + """Return (organism_record, files) for a MycoCosm portal id (e.g. "Acain1") + from the JGI Data Portal search endpoint, paginating to gather every file. + Returns (None, []) if the organism is absent or the query fails.""" + org, files, page = None, [], 1 + while True: + params = {"organism": portal_id, "api_version": "2", "x": "50", "p": str(page)} + data = None + for attempt in range(1, max_attempts + 1): + try: + resp = session.get( + SEARCH_URL, + params=params, + headers={"accept": "application/json"}, + timeout=120, + ) + except requests.RequestException as error: + # transient per-attempt search noise during bulk assimilation; + # DEBUG so it is hidden unless --verbose is set + logger.debug( + f"{spacer}\t{portal_id} search error (attempt {attempt}): {error}" + ) + time.sleep(2) + continue + if resp.status_code != 200: + logger.debug(f"{spacer}\t{portal_id} search HTTP {resp.status_code}") + time.sleep(2) + continue + try: + data = resp.json() + except ValueError: + logger.debug(f"{spacer}\t{portal_id} search returned non-JSON") + time.sleep(2) + continue + break + if data is None: + return None, [] + organisms = data.get("organisms") or [] + if not organisms: + break + if org is None: + org = organisms[0] + files.extend(organisms[0].get("files") or []) + total = data.get("file_total") or len(files) + if len(files) >= total or not data.get("next_page"): + break + page += 1 + return org, files + + +def _mycocosm_ids(org_id, top_hit, portal_id, file_ids): + """Build the request_archived_files / download_files ``ids`` payload for a + MycoCosm organism.""" + entry = {"file_ids": list(file_ids)} + if top_hit: + entry["top_hit"] = top_hit + if portal_id: + entry["mycocosm_portal_id"] = portal_id + return {org_id: entry} + + +def request_restore(session, token, ids_payload, spacer="\t"): + """Request that archived (PURGED) files be restored to disk. Returns the + restore request's status URL, or None on failure.""" + body = {"ids": ids_payload, "send_mail": False, "api_version": "2"} + try: + resp = session.post( + RESTORE_URL, + json=body, + headers={ + "accept": "application/json", + "content-type": "application/json", + "Authorization": "Bearer " + token, + }, + timeout=120, + ) + except requests.RequestException as error: + logger.warning(f"{spacer}\trestore request error: {error}") + return None + if resp.status_code != 200: + logger.warning(f"{spacer}\trestore request failed (HTTP {resp.status_code})") + return None + try: + return resp.json().get("request_status_url") + except ValueError: + return None + + +def _fmt_elapsed(seconds): + """Human-friendly mm:ss / h:mm elapsed string.""" + seconds = int(seconds) + if seconds < 3600: + return f"{seconds // 60}m{seconds % 60:02d}s" + return f"{seconds // 3600}h{(seconds % 3600) // 60:02d}m" + + +# what each JGI restore status means for a tape->disk transfer, surfaced to users +_RESTORE_STATUS_MSG = { + "new": "request queued", + "pending": "retrieving from tape", + "staging": "staging to disk", + "ready": "staged to disk", + "expired": "restore expired", +} + + +def poll_restore( + session, status_url, timeout=60, interval=30, spacer="\t", label="", heartbeat=60 +): + """Poll a tape-restore request until its files are READY (returns True) or + the timeout / expiry is reached (returns False). A `timeout` of None polls + until the restore resolves one way or the other, however long that takes. + + Progress is logged so the user can see the tape->disk transfer advance: + every status transition (queued -> retrieving -> staging -> ready) is + reported, plus a heartbeat every `heartbeat` seconds while a stage lingers.""" + if not status_url: + return False + tag = f"{label}: " if label else "" + waited = 0 + last_status = None + last_heartbeat = 0 + while timeout is None or waited <= timeout: + status = "" + try: + resp = session.get( + status_url, headers={"accept": "application/json"}, timeout=60 + ) + status = (resp.json().get("status") or "").lower() + except (requests.RequestException, ValueError): + pass + if status == "ready": + logger.info( + f"{spacer}\t{tag}tape restore complete - files staged to disk " + f"(waited {_fmt_elapsed(waited)})" + ) + return True + if status == "expired": + logger.warning( + f"{spacer}\t{tag}tape restore expired; a new request is needed" + ) + return False + # surface the transfer's progress: log each stage change, then a + # periodic heartbeat so a long-running stage does not look hung + detail = _RESTORE_STATUS_MSG.get(status, status or "waiting") + if status != last_status: + logger.info( + f"{spacer}\t{tag}tape restore: {detail} " + f"(elapsed {_fmt_elapsed(waited)}; disk restores usually take " + "<1 h, up to a night)" + ) + last_status = status + last_heartbeat = waited + elif heartbeat and waited - last_heartbeat >= heartbeat: + logger.info( + f"{spacer}\t{tag}tape restore still in progress: {detail} " + f"(elapsed {_fmt_elapsed(waited)})" + ) + last_heartbeat = waited + time.sleep(interval) + waited += interval + logger.warning( + f"{spacer}\t{tag}tape restore did not complete within {_fmt_elapsed(timeout)}" + ) + return False + + +def download_zip(session, token, ids_payload, dest_zip, spacer="\t", max_attempts=3): + """Download the given RESTORED files as a single zip stream. Returns True on + success (dest_zip written), False otherwise.""" + body = {"ids": ids_payload, "api_version": "2"} + for attempt in range(1, max_attempts + 1): + try: + resp = session.post( + DOWNLOAD_URL, + json=body, + headers={ + "accept": "application/json", + "content-type": "application/json", + "Authorization": "Bearer " + token, + }, + timeout=1800, + stream=True, + ) + except requests.RequestException as error: + logger.warning(f"{spacer}\tdownload error (attempt {attempt}): {error}") + time.sleep(5) + continue + ctype = resp.headers.get("content-type", "") + if resp.status_code == 200 and "zip" in ctype.lower(): + with open(dest_zip, "wb") as out: + for chunk in resp.iter_content(chunk_size=1 << 20): + if chunk: + out.write(chunk) + resp.close() + return True + resp.close() + logger.warning( + f"{spacer}\tdownload attempt {attempt} failed (HTTP {resp.status_code}, {ctype})" + ) + time.sleep(5) + return False + + +def extract_zip(zip_path, wanted, spacer="\t"): + """Extract files from a JGI download archive. `wanted` maps a member's + basename -> destination path. Returns the set of basenames extracted.""" + extracted = set() + try: + with zipfile.ZipFile(zip_path) as archive: + members = {Path(m).name: m for m in archive.namelist()} + for basename, dest in wanted.items(): + member = members.get(basename) + if member is None: + continue + Path(dest).parent.mkdir(parents=True, exist_ok=True) + with archive.open(member) as src, open(dest, "wb") as out: + shutil.copyfileobj(src, out) + extracted.add(basename) + except zipfile.BadZipFile: + logger.error(f"{spacer}\tcorrupt JGI download archive {zip_path}") + return extracted + + +def _dwnld_org( + session, token, ids_payload, portal_id, selected, output, tmp_dir, spacer +): + """Download one organism's `selected` {type: file record} as a single zip and + extract it into `output//`. Returns {type: destination path} for the + files actually obtained (empty when the download itself failed).""" + dest_zip = os.path.join(tmp_dir, f"{portal_id}.zip") + if not download_zip(session, token, ids_payload, dest_zip, spacer=spacer): + logger.warning(f"{spacer}\t{portal_id}: download failed") + return {} + + wanted, type_dest = {}, {} + for typ, f in selected.items(): + name = f["file_name"] + dest = os.path.join(output, typ, name) + wanted[name] = dest + type_dest[typ] = (name, dest) + extracted = extract_zip(dest_zip, wanted, spacer=spacer) + if Path(dest_zip).is_file(): + Path(dest_zip).unlink() + + obtained = {} + for typ, (name, dest) in type_dest.items(): + if name in extracted and Path(dest).is_file(): + obtained[typ] = dest + logger.info(f"{spacer}\t{portal_id} {typ}: {name}") + else: + logger.warning(f"{spacer}\t{portal_id}: {typ} missing from archive") + return obtained + + +def _flush_restores(session, token, buffer, spacer="\t"): + """Request restores for a batch of organisms in one call - the + request_archived_files ``ids`` payload is keyed by organism, so many + organisms ride on a single request - and stamp each with the time its + restore was asked for, which starts its wait clock.""" + if not buffer: + return + ids_payload = {} + for entry in buffer: + ids_payload.update(entry["ids"]) + request_restore(session, token, ids_payload, spacer=spacer) + requested_at = time.time() + for entry in buffer: + entry["requested_at"] = requested_at + entry["last_request"] = requested_at + logger.info( + f"{spacer}\trequested tape restores for {len(buffer)} organism(s): " + + ", ".join(e["portal_id"] for e in buffer[:5]) + + (", ..." if len(buffer) > 5 else "") + ) + buffer.clear() + + +# a restore request JGI drops or lets expire would otherwise strand an +# indefinite wait forever, so a still-pending organism is re-requested this +# often (seconds) +_REREQUEST_RESTORE = 6 * 3600 + + +def _fmt_wait(minutes): + """Human-friendly rendering of a wait allowance given in minutes; None is an + unbounded wait.""" + if minutes is None: + return "as long as it takes" + minutes = int(minutes) + if minutes < 60: + return f"{minutes}m" + if minutes % 60: + return f"{minutes // 60}h{minutes % 60:02d}m" + return f"{minutes // 60}h" + + +def circle_back( + session, + token, + pending, + df, + output, + tmp_dir, + dwnlds, + ome_set, + deferred, + masked=True, + restore_wait=None, + poll_interval=60, + spacer="\t", +): + """Revisit organisms whose files were left on tape during the main pass, + downloading each as JGI stages it to disk. + + Availability is re-read from ``mycocosm_file_list`` per organism - that is + the authoritative signal. The restore request's own status URL lags behind + the files it restored (it still reports `pending` once they are RESTORED), + and the generic ``/search/?datasets=`` endpoint, though it accepts many + organisms at once, silently omits some MycoCosm portals; neither can be + trusted here. + + One sweep of the pending organisms is made per `poll_interval` seconds. An + organism is downloaded as soon as every file it still needs is RESTORED. + `restore_wait` is how many minutes any one organism is waited on before it + is given up on - deferred, not failed; None (the default) waits for as long + as JGI takes, re-requesting a restore that has gone stale.""" + wait_seconds = None if restore_wait is None else max(0, restore_wait * 60) + logger.info( + f"{spacer}{len(pending)} genome(s) awaiting tape restore; checking every " + f"{poll_interval}s, waiting {_fmt_wait(restore_wait)} per genome" + ) + started = time.time() + while pending: + sweep_start = time.time() + for portal_id in list(pending): + entry = pending[portal_id] + requested_at = entry.get("requested_at") or sweep_start + waited = time.time() - requested_at + + # a dropped or expired restore request would strand an unbounded + # wait, so reissue one that has been outstanding too long + if time.time() - (entry.get("last_request") or requested_at) >= ( + _REREQUEST_RESTORE + ): + logger.info( + f"{spacer}\t{portal_id}: still on tape after " + f"{_fmt_elapsed(waited)}; re-requesting its restore" + ) + request_restore(session, token, entry["ids"], spacer=spacer) + entry["last_request"] = time.time() + + org, files = search_organism(session, portal_id, spacer=spacer) + still_needed = {} + if org: + for typ in entry["typs"]: + chosen = select_file(files, typ, masked=masked) + if chosen is not None: + still_needed[typ] = chosen + if not still_needed: + # JGI publishes none of these files any more - a permanent + # condition, so do not spend the wait allowance on it + logger.warning( + f"{spacer}\t{portal_id}: JGI no longer lists the requested " + "file(s)" + ) + _finish_org( + df, + entry["i"], + portal_id, + dwnlds, + entry["dwnlded"], + ome_set, + spacer, + ) + del pending[portal_id] + continue + + if still_needed and all(_is_restored(f) for f in still_needed.values()): + logger.info( + f"{spacer}\t{portal_id}: staged to disk after " + f"{_fmt_elapsed(waited)}; downloading" + ) + ids_payload = _mycocosm_ids( + org.get("id"), + (org.get("top_hit") or {}).get("_id"), + org.get("mycocosm_portal_id") or portal_id, + [f["_id"] for f in still_needed.values()], + ) + entry["dwnlded"].update( + _dwnld_org( + session, + token, + ids_payload, + portal_id, + still_needed, + output, + tmp_dir, + spacer, + ) + ) + _finish_org( + df, entry["i"], portal_id, dwnlds, entry["dwnlded"], ome_set, spacer + ) + del pending[portal_id] + elif wait_seconds is not None and waited >= wait_seconds: + logger.warning( + f"{spacer}\t{portal_id}: still on tape after " + f"{_fmt_elapsed(waited)}; deferring to a later run" + ) + ome_set.add(portal_id) + deferred.add(portal_id) + del pending[portal_id] + + # spread a sweep's status checks over the interval so a large + # backlog never bursts requests at JGI + if pending: + time.sleep(min(poll_interval / len(pending), 1)) + + if pending: + # an unbounded wait can be a long one; say what it is waiting on + logger.info( + f"{spacer}\t{len(pending)} genome(s) still on tape after " + f"{_fmt_elapsed(time.time() - started)}" + ) + remaining = poll_interval - (time.time() - sweep_start) + if remaining > 0: + time.sleep(remaining) + + +def _finish_org(df, i, portal_id, dwnlds, dwnlded, ome_set, spacer="\t"): + """Record an organism's retrieved files as ``_path`` columns and flag + it as failed when an essential type (assembly or gff3) was not obtained - + whether JGI had no such file, or the download/extraction did not yield it.""" + for typ, path in dwnlded.items(): + df.at[i, typ + "_path"] = path + essential = [t for t in ("fna", "gff3") if t in dwnlds and t not in dwnlded] + if essential: + logger.warning( + f"{spacer}\t{portal_id}: no {'/'.join(essential)} retrieved; excluding" + ) + ome_set.add(portal_id) + + +def main( + df, + output, + user, + pwd, + assembly=True, + proteome=False, + gff3=True, + transcript=False, + est=False, + masked=True, + spacer="\t", + restore_wait=None, + poll_interval=60, + request_delay=3, + ome_col=None, + deferred=None, + defer_tape=False, + restore_chunk=50, +): + """Download MycoCosm data for the JGI portal ids in `df` via the JGI Data + Portal API (non-Globus): `df` gains + ``_path`` columns (e.g. fna_path, gff3_path) plus genus/species/strain, + and the function returns (df, failed_portal_ids). + + `ome_col` names the portal id column; it defaults to ``assembly_acc``, or the + lone column of a single-column input (MycoCosm tables label it ``portal``). + + Most MycoCosm files are archived on tape (file_status PURGED) and must be + staged to disk before they can be downloaded. `defer_tape` chooses how that + wait is spent: + + - False (default): each organism's restore is awaited in place, up to + `restore_wait` minutes, before moving to the next portal id. + - True: an organism needing tape I/O is skipped immediately - its restore + is requested (batched `restore_chunk` organisms to a call) and it is + revisited only after every portal id has been visited, then polled every + `poll_interval` seconds until staged. Far faster over a large table, + where nearly every genome needs a restore. + + Either way `restore_wait` is the maximum a single genome is waited on, in + minutes, after which it is deferred rather than failed. None (the default) + waits for as long as JGI takes. Pass a set as `deferred` to receive the + portal ids given up on: unlike genuine failures (portal absent from + MycoCosm, no such file type, corrupt download) they are pending JGI tape + I/O, so callers should retry them rather than blacklist them.""" + if deferred is None: + deferred = set() + if ome_col is not None: + if ome_col not in df.columns: + logger.error(f"Invalid input. No {ome_col} column.") + return df, set() + elif "assembly_acc" in df.columns: + ome_col = "assembly_acc" + elif len(df.columns) == 1: + ome_col = list(df.columns)[0] + else: + logger.error("Invalid input. No assembly_acc column and more than one column.") + return df, set() + + dwnlds = [] + if assembly: + dwnlds.append("fna") + if proteome: + dwnlds.append("faa") + if gff3: + dwnlds.append("gff3") + if transcript: + dwnlds.append("transcript") + if est: + dwnlds.append("est") + + output = str(output).rstrip("/") + for typ in dwnlds: + Path(os.path.join(output, typ)).mkdir(parents=True, exist_ok=True) + tmp_dir = os.path.join(output, "jgi_zip") + Path(tmp_dir).mkdir(parents=True, exist_ok=True) + + logger.info(spacer + "Logging into JGI") + session, token = jgi_api_login(user, pwd, spacer=spacer) + + logger.info( + f"{spacer}Downloading {len(df)} JGI organism(s) via the JGI Data Portal API" + ) + ome_set = set() + pending, restore_buffer = {}, [] + for i, row in tqdm(df.iterrows(), total=len(df)): + portal_id = row[ome_col] + + org, files = search_organism(session, portal_id, spacer=spacer) + if not org: + # expected during bulk assimilation; DEBUG so it is hidden unless + # --verbose is set + logger.debug(f"{spacer}\t{portal_id} not found in JGI MycoCosm") + ome_set.add(portal_id) + continue + + org_id = org.get("id") + top_hit = (org.get("top_hit") or {}).get("_id") + portal = org.get("mycocosm_portal_id") or portal_id + + genus, species, strain = parse_org_name(org.get("name") or "") + _fill_if_empty(df, i, "genus", genus) + _fill_if_empty(df, i, "species", species) + _fill_if_empty(df, i, "strain", strain) + + selected, absent = {}, [] + for typ in dwnlds: + chosen = select_file(files, typ, masked=masked) + if chosen is None: + absent.append(typ) + else: + selected[typ] = chosen + if absent: + logger.warning(f"{spacer}\t{portal_id}: JGI has no {'/'.join(absent)} file") + + # resume a previous run: any file already retrieved is kept as-is. + # Downloads arrive gzipped, but curation decompresses in place, so an + # unzipped copy counts too + dwnlded = {} + for typ, f in tuple(selected.items()): + dest = os.path.join(output, typ, f["file_name"]) + for path in (dest, re.sub(r"\.gz$", "", dest)): + if Path(path).is_file() and Path(path).stat().st_size > 0: + dwnlded[typ] = path + del selected[typ] + logger.info( + f"{spacer}\t{portal_id} {typ}: {Path(path).name} (preexisting)" + ) + break + + if not selected: + # nothing left to retrieve - either all preexisting, or JGI has no + # files for any requested type + _finish_org(df, i, portal_id, dwnlds, dwnlded, ome_set, spacer) + continue + + # JGI keeps most files in tape archive (file_status PURGED); those must + # be transferred to disk (RESTORED) before they can be downloaded. Report + # which files are on tape vs already on disk, then request the restore. + on_disk = [f for f in selected.values() if _is_restored(f)] + on_tape = [f for f in selected.values() if not _is_restored(f)] + if on_disk: + logger.info( + f"{spacer}\t{portal_id}: {len(on_disk)} file(s) already on disk: " + + ", ".join(f.get("file_name", f["_id"]) for f in on_disk) + ) + if on_tape: + logger.debug( + f"{spacer}\t{portal_id}: {len(on_tape)} file(s) are archived on TAPE and " + "must be restored to disk before download - " + + ", ".join(f.get("file_name", f["_id"]) for f in on_tape) + ) + tape_ids = _mycocosm_ids( + org_id, top_hit, portal, [f["_id"] for f in on_tape] + ) + if defer_tape: + # do not block the pass on JGI tape I/O: queue the restore and + # come back to this organism once every portal has been visited + pending[portal_id] = { + "i": i, + "portal_id": portal_id, + "typs": list(selected), + "dwnlded": dwnlded, + "ids": tape_ids, + "requested_at": None, + } + restore_buffer.append(pending[portal_id]) + if len(restore_buffer) >= restore_chunk: + _flush_restores(session, token, restore_buffer, spacer=spacer) + continue + + status_url = request_restore(session, token, tape_ids, spacer=spacer) + if not poll_restore( + session, + status_url, + timeout=None if restore_wait is None else restore_wait * 60, + interval=poll_interval, + spacer=spacer, + label=portal_id, + ): + logger.warning( + f"{spacer}\t{portal_id}: tape restore still pending; deferring this " + "genome (rerun later to resume once JGI has staged it to disk)" + ) + ome_set.add(portal_id) + deferred.add(portal_id) + continue + + ids_payload = _mycocosm_ids( + org_id, top_hit, portal, [f["_id"] for f in selected.values()] + ) + dwnlded.update( + _dwnld_org( + session, + token, + ids_payload, + portal_id, + selected, + output, + tmp_dir, + spacer, + ) + ) + _finish_org(df, i, portal_id, dwnlds, dwnlded, ome_set, spacer) + if request_delay: + time.sleep(request_delay) + + # circle back to the organisms skipped for tape restores above + _flush_restores(session, token, restore_buffer, spacer=spacer) + if pending: + circle_back( + session, + token, + pending, + df, + output, + tmp_dir, + dwnlds, + ome_set, + deferred, + masked=masked, + restore_wait=restore_wait, + poll_interval=poll_interval, + spacer=spacer, + ) + + # tidy the scratch zip dir if empty + try: + Path(tmp_dir).rmdir() + except OSError: + pass + + for col in ("gff3", "faa", "fna"): + if col in df.columns: + del df[col] + + return df, ome_set + + +def cli(): + + parser = argparse.ArgumentParser( + description="Imports table/database with a JGI `assembly_acc` column (MycoCosm " + + "portal ids) and downloads assembly, proteome, and/or gff3 via the JGI Data " + + "Portal API. Supports rerunning/continuing previous runs in the same directory." + ) + parser.add_argument( + "-i", + "--input", + required=True, + help="Genome code or table with `assembly_acc` column of JGI ome codes", + ) + parser.add_argument( + "-a", + "--assembly", + default=False, + action="store_true", + help="Download assembly fastas", + ) + parser.add_argument( + "-p", + "--proteome", + default=False, + action="store_true", + help="Download proteome fastas", + ) + parser.add_argument( + "-g", "--gff", default=False, action="store_true", help="Download gff3s" + ) + parser.add_argument( + "-t", + "--transcript", + default=False, + action="store_true", + help="Download transcripts fastas", + ) + parser.add_argument( + "-e", "--est", default=False, action="store_true", help="Download EST fastas" + ) + parser.add_argument( + "--nonmasked", + default=False, + action="store_true", + help="[-a] Download nonmasked assemblies", + ) + parser.add_argument( + "-s", + "--skip-tape", + default=False, + action="store_true", + help="Skip genomes with tape-archived files, request their restore, and " + + "circle back to download them once JGI stages them to disk", + ) + parser.add_argument( + "-w", + "--tape-wait", + type=int, + default=None, + help="Maximum minutes to wait for a single genome's tape restore before " + + "deferring it to a later run. DEFAULT: wait indefinitely", + ) + parser.add_argument("-o", "--output", default=str(Path.cwd()), help="Output dir") + parser.add_argument( + "-v", + "--verbose", + default=False, + action="store_true", + help="Report per-genome search/download diagnostics (DEBUG logging)", + ) + args = parser.parse_args() + setup_logging(verbose=getattr(args, "verbose", False)) + + if args.nonmasked: + args.assembly = True + + if ( + not args.assembly + and not args.proteome + and not args.transcript + and not args.est + and not args.gff + ): + logger.error("You must choose at least one download option.") + + ncbi_api, user, pwd = login_check(ncbi=False) + + args_dict = { + "JGI Table": args.input, + "Assemblies": args.assembly, + "RepeatMasked": not args.nonmasked, + "Proteomes": args.proteome, + ".gff3's": args.gff, + "Transcripts": args.transcript, + "EST": args.est, + "Skip tape files": args.skip_tape, + "Max tape wait": ( + "indefinite" if args.tape_wait is None else f"{args.tape_wait} minute(s)" + ), + } + + start_time = intro("Download JGI files", args_dict) + logger.warning( + "This script does NOT account for use-restricted data. " + + "It is user responsibility to determine use restriction status " + + "in accord with the MycoCosm terms and conditions: " + + "https://jgi.doe.gov/user-programs/pmo-overview/policies/legacy-data-policies/" + ) + logger.info("") + + if Path(args.input).is_file(): + with open(args.input, "r") as raw: + for line in raw: + if "assembly_acc" in line.rstrip().split("\t"): + df = pd.read_csv(args.input, sep="\t", index_col=None) + else: + df = pd.read_csv(args.input, sep="\t", header=None) + break + else: + in_data = args.input.replace('"', "").replace("'", "").replace(",", " ").split() + df = pd.DataFrame({"assembly_acc": in_data}) + + output = format_path(args.output) + + jgi_df, ome_set = main( + df, + output, + user, + pwd, + args.assembly, + args.proteome, + args.gff, + args.transcript, + args.est, + not args.nonmasked, + spacer="", + restore_wait=args.tape_wait, + defer_tape=args.skip_tape, + ) + jgi_df = jgi_df.rename(columns={"assembly_acc": "#assembly_acc"}) + jgi_df["source"] = "jgi" + jgi_df["restriction"] = "no" + jgi_df.to_csv(str(Path(args.input)) + ".predb.tsv", sep="\t", index=False) + + outro(start_time) + + +if __name__ == "__main__": + cli() diff --git a/mycotools/ncbiDwnld.py b/mycotools/download/ncbi.py similarity index 56% rename from mycotools/ncbiDwnld.py rename to mycotools/download/ncbi.py index 826f2d3..fffe7fd 100755 --- a/mycotools/ncbiDwnld.py +++ b/mycotools/download/ncbi.py @@ -1,23 +1,22 @@ #! /usr/bin/env python3 # NEED a db check to ensure the log is relevant to the input -# NEED to convert to datasets # NEED to consider refseq genomes with annotations when genbank doesn't have them import os import re import sys -import gzip import json +import math import time import shutil import urllib +import logging import zipfile import argparse +import warnings import subprocess -import numpy as np import pandas as pd -from contextlib import closing from tqdm import tqdm from Bio import Entrez from datetime import datetime @@ -25,23 +24,23 @@ intro, outro, format_path, - prep_output, - mkOutput, - eprint, - vprint, - findExecs, + mk_output, + find_execs, read_json, split_input, + setup_logging, ) -from mycotools.lib.dbtools import log_editor, loginCheck, mtdb, read_tax +from mycotools.lib.dbtools import clean_api_key, log_editor, login_check, mtdb +from pathlib import Path + +logger = logging.getLogger(__name__) pd.options.mode.chained_assignment = None def ncbidb2df(data, stdin=False): - import pandas as pd, pandas + import pandas as pd - columns = mtdb.columns if isinstance(data, mtdb): db_df = pd.DataFrame(data.reset_index()) elif not stdin: @@ -63,20 +62,20 @@ def prepare_folders(output_path, gff, prot, assem, transcript): file_types = [] if assem: - if not os.path.exists(output_path + "fna"): - os.mkdir(output_path + "fna") + if not Path(output_path + "fna").exists(): + Path(output_path + "fna").mkdir() file_types.append("fna") if gff: - if not os.path.exists(output_path + "gff3"): - os.mkdir(output_path + "gff3") + if not Path(output_path + "gff3").exists(): + Path(output_path + "gff3").mkdir() file_types.append("gff3") if prot: - if not os.path.exists(output_path + "faa"): - os.mkdir(output_path + "faa") + if not Path(output_path + "faa").exists(): + Path(output_path + "faa").mkdir() file_types.append("faa") if transcript: - if not os.path.exists(output_path + "transcript"): - os.mkdir(output_path + "transcript") + if not Path(output_path + "transcript").exists(): + Path(output_path + "transcript").mkdir() file_types.append("transcript") return file_types @@ -85,7 +84,7 @@ def prepare_folders(output_path, gff, prot, assem, transcript): def compile_log(output_path): acc2log = {} - if not os.path.isfile(output_path): + if not Path(output_path).is_file(): with open(output_path, "w") as out: out.write("#acc\tassembly_acc\n") else: @@ -116,11 +115,11 @@ def esearch_ncbi(accession, column, database="assembly"): handle = Entrez.esearch(db=database, term=search_term) genome_ids = Entrez.read(handle)["IdList"] break - except (RuntimeError, urllib.error.HTTPError) as e: + except (RuntimeError, urllib.error.HTTPError): time.sleep(1) esc_count += 1 else: - print("\tERROR:", accession, "failed to search NCBI") + logger.error(f"{accession} failed to search NCBI") return None return genome_ids @@ -138,7 +137,7 @@ def esummary_ncbi(ID, database): continue if database == "assembly": try: # is it populated with an FTP? - ftp_path = str( + str( record["DocumentSummarySet"]["DocumentSummary"][0][ "FtpPath_GenBank" ] @@ -214,7 +213,7 @@ def collect_assembly_accs( if not genome_id: # No IDs retrieved if "ome" in row.keys(): accession = row["ome"] - eprint(spacer + "\t" + accession + " failed to find genome ID", flush=True) + logger.error(spacer + "\t" + accession + " failed to find genome ID") try: failed.append([accession, datetime.strftime(row["version"], "%Y%m%d")]) except TypeError: # if the row can't be formatted as a date @@ -241,7 +240,6 @@ def collect_assembly_accs( record_info = record["DocumentSummarySet"]["DocumentSummary"][0] assemblyID = record_info["AssemblyAccession"] - esc_count = 0 log_editor( output_path + "ncbiDwnld.log", str(new_acc), @@ -257,8 +255,22 @@ def collect_assembly_accs( return acc2log, failed, out_df -def run_datasets(include, accs_file, output_path, annotated, api=None, verbose=False): - """Run NCBI datasets to download genomes or metadata""" +def run_datasets( + include, + accs_file, + output_path, + annotated, + api=None, + verbose=False, + mute_stderr=False, +): + """Run NCBI datasets to download genomes or metadata + + `mute_stderr` keeps datasets' own progress bar and error text off the + terminal entirely; it is captured and logged at debug instead. It exists for + callers that draw their own progress bar and would otherwise be overdrawn, + so it belongs to the caller rather than the CLI and is not exposed as a + flag. The caller stays responsible for reporting the failure itself.""" dataset_scaf = [ "datasets", "download", @@ -271,18 +283,36 @@ def run_datasets(include, accs_file, output_path, annotated, api=None, verbose=F dataset_scaf.extend(["--include", include]) else: dataset_scaf.append("--dehydrated") + api = clean_api_key(api) if api: dataset_scaf += ["--api-key", api] if annotated: dataset_scaf.append("--annotated") - cwd = os.getcwd() + cwd = str(Path.cwd()) os.chdir(output_path) if verbose: - v = None + dataset_call = subprocess.call(dataset_scaf) else: - v = subprocess.DEVNULL - dataset_call = subprocess.call(dataset_scaf, stdout=v, stderr=v) + # capture stderr rather than discard it: datasets reports why it failed + # there, and a silenced call that fails leaves no other explanation + proc = subprocess.run( + dataset_scaf, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, text=True + ) + dataset_call = proc.returncode + # datasets' progress bar shares this stream, so collapse each line to + # its last carriage-return frame rather than reprinting every redraw + # (note: splitlines() would break on \r, hence the explicit split) + err = "\n".join( + x.split("\r")[-1].rstrip() + for x in proc.stderr.split("\n") + if x.split("\r")[-1].strip() + ) + if err: + # the bar lands here too, so it is only worth surfacing on failure + # -- and not even then once a caller has muted it + loud = dataset_call and not mute_stderr + (logger.error if loud else logger.debug)(err) os.chdir(cwd) return dataset_call @@ -344,15 +374,87 @@ def compile_organism_names(unzip_path, spacer="\t"): return acc2org, acc2meta, failed +def resolve_paired_accs( + accs, acc_file, output_path, api=None, spacer="\t\t", summary_chunk=500 +): + """Map accessions onto their counterpart in the other NCBI repository. + + GenBank and RefSeq version their assemblies independently, so the + counterpart of GCA_017499595.2 is GCF_017499595.1 -- swapping the prefix + and keeping the version names GCF_017499595.2, which does not exist. NCBI + reports the pair it actually holds, and reports none for an assembly that + was never mirrored, so an accession missing from the result has nothing to + reattempt rather than a counterpart that failed. + + These are metadata records rather than genomes, so the chunk is sized well + above the genome chunk for the same reason the data reports are. A chunk + that cannot be resolved is skipped rather than fatal: the accessions in it + simply keep their existing failure, which is the outcome without a + counterpart anyway.""" + + acc2pair = {} + accs = [str(x) for x in accs] + for i in range(0, len(accs), summary_chunk): + batch = accs[i : i + summary_chunk] + with open(acc_file, "w") as out: + out.write("\n".join(batch)) + cmd = [ + "datasets", + "summary", + "genome", + "accession", + "--inputfile", + acc_file, + "--as-json-lines", + ] + api_key = clean_api_key(api) + if api_key: + cmd += ["--api-key", api_key] + + cwd = str(Path.cwd()) + os.chdir(output_path) + try: + proc = subprocess.run( + cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True + ) + # an absent datasets is the caller's problem to report, not a reason to + # end a run that has already downloaded everything it could + except FileNotFoundError: + logger.debug(f"{spacer}\tdatasets is unavailable to resolve pairs") + os.chdir(cwd) + return acc2pair + os.chdir(cwd) + if proc.returncode: + logger.debug( + f"{spacer}\tcould not resolve paired accessions: " + + f"{proc.stderr.rstrip()}" + ) + continue + + for line in proc.stdout.split("\n"): + if not line.strip(): + continue + try: + report = json.loads(line) + except json.JSONDecodeError: + continue + if report.get("accession") and report.get("paired_accession"): + acc2pair[report["accession"]] = report["paired_accession"] + + return acc2pair + + def parse_datasets(datasets_path, unzip_base, req_files, spacer="\t"): """Unzip, identify complete downloads, parse file outputs and metadata, report missing data to check alternative repository""" try: with zipfile.ZipFile(datasets_path, "r") as zip_ref: zip_ref.extractall(unzip_base) - except zipfile.BadZipFile: + # a dropped transfer leaves a truncated archive, and a call that died before + # it opened the file leaves none at all; both are the same failure here + except (zipfile.BadZipFile, FileNotFoundError): return False, False, False - os.remove(datasets_path) + Path(datasets_path).unlink() unzip_path = unzip_base + "ncbi_dataset/" type2ncbi = { @@ -376,8 +478,8 @@ def parse_datasets(datasets_path, unzip_base, req_files, spacer="\t"): if req_files.difference(set(files.keys())): failed.append(data["accession"]) for t, f_ in files.items(): - if os.path.isfile(f_): - os.remove(f_) + if Path(f_).is_file(): + Path(f_).unlink() else: acc2data[data["accession"]] = files @@ -386,6 +488,123 @@ def parse_datasets(datasets_path, unzip_base, req_files, spacer="\t"): return acc2data, acc2org, failed +def download_datasets( + accs, + acc_file, + include, + req_files, + annotated, + output_path, + api=None, + verbose=False, + spacer="\t\t", + max_attempts=3, + min_accs=10, + max_dead=None, +): + """Write a chunk of accessions to acc_file, download them via NCBI datasets, + and parse the output. Returns acc2files, acc2org, failed (each False only if + nothing in `accs` could be retrieved). + + NCBI drops these transfers mid-stream, and the larger the archive the likelier + it is to be dropped -- so a batch that fails every attempt is halved and its + halves are downloaded separately rather than retried whole. Retrying whole + restarts a multi-GB transfer from zero at the size that was already failing + and loses the entire batch when it fails again; halving both shrinks the + transfer and keeps whatever the other half retrieved. + + Splitting is abandoned once `max_dead` batches have come back with nothing + and nothing at all has been retrieved, because a repository serving nothing + at any size is down rather than overloaded and halving it only multiplies the + calls made against it. The default tolerance is the depth of the split tree: + halving descends the first half before trying its sibling, so the first batch + that can succeed may be that many failures away.""" + zip_path = output_path + "ncbi_dataset.zip" + if max_dead is None: + max_dead = 2 + math.ceil(math.log2(max(len(accs) / max(min_accs, 1), 1))) + + def attempt_batch(batch): + """Download one batch, retrying a dropped transfer at the same size. + Returns the parsed result, or None if every attempt failed.""" + for attempt in range(1, max_attempts + 1): + if attempt > 1: + # datasets leaves a truncated archive behind when the stream is + # reset, and a call that dies before reopening it would hand + # that same partial file back to the parser + Path(zip_path).unlink(missing_ok=True) + # a reset usually means NCBI is loaded; retrying instantly adds + # to the load that caused it + time.sleep(3 * 2 ** (attempt - 2)) + logger.debug(f"{spacer}\tAttempt {attempt} ({len(batch)} accessions)") + + with open(acc_file, "w") as out: + out.write("\n".join([str(x) for x in batch])) + + code = run_datasets( + include, + acc_file, + output_path, + api=api, + verbose=verbose, + annotated=annotated, + # a failed attempt is routine now that it is retried and split, + # so datasets' error text is debug detail; the failures that + # survive the splitting are what the caller is told about + mute_stderr=True, + ) + if code: + continue + + parsed = parse_datasets(zip_path, output_path, req_files, spacer) + if parsed[0] is not False: + return parsed + Path(zip_path).unlink(missing_ok=True) + return None + + logger.debug(f"{spacer}Downloading data") + acc2files, acc2org, failed = {}, {}, [] + queue, retrieved, dead_streak = [list(accs)], False, 0 + while queue: + batch = queue.pop(0) + parsed = attempt_batch(batch) + if parsed is not None: + retrieved = True + dead_streak = 0 + acc2files.update(parsed[0]) + acc2org.update(parsed[1]) + failed.extend(parsed[2]) + continue + + # once anything has come back the repository is serving, and every later + # failure is that batch's problem rather than grounds to stop splitting + if not retrieved: + dead_streak += 1 + if dead_streak >= max_dead: + logger.debug( + f"{spacer}\t{dead_streak} batches returned nothing; abandoning " + + f"{sum(len(b) for b in queue) + len(batch)} accessions" + ) + break + # below this size the transfer is no longer what is failing, so the + # accessions are abandoned to the caller rather than split again. It + # derives what it never received from acc2files, so they need no + # accounting here + if len(batch) <= min_accs: + continue + + mid = len(batch) // 2 + logger.debug( + f"{spacer}\t{len(batch)} accessions failed {max_attempts} attempts; " + + "splitting" + ) + queue[:0] = [batch[:mid], batch[mid:]] + + if not retrieved: + return False, False, False + + return acc2files, acc2org, failed + + def main( api=None, assembly=True, @@ -394,12 +613,13 @@ def main( transcript=False, ncbi_df=False, remove=False, - output_path=os.getcwd(), + output_path=str(Path.cwd()), verbose=False, column="assembly_acc", ncbi_column="Assembly", check_MD5=True, spacer="\t\t", + chunk=25, ): # initialize run directory and information @@ -408,7 +628,7 @@ def main( # assembly, transcript) # check if ncbi_df is a dataframe, and import if not - if not isinstance(ncbi_df, pd.DataFrame) and os.path.isfile(ncbi_df): + if not isinstance(ncbi_df, pd.DataFrame) and Path(ncbi_df).is_file(): ncbi_df = ncbidb2df(ncbi_df) if len(ncbi_df.index) == 0: ncbi_df = pd.DataFrame({i: [v] for i, v in enumerate(list(ncbi_df.keys()))}) @@ -420,12 +640,19 @@ def main( elif "version" not in ncbi_df.keys(): ncbi_df["version"] = "" + # Uppercase before the index is taken from it, not after. The index is what + # downloads and failures are matched on downstream, and NCBI reports its + # accessions uppercase, so normalizing the column afterwards leaves a + # lowercase index that matches neither and drops the row out of the run + if "assembly_acc" in ncbi_df.keys(): + ncbi_df["assembly_acc"] = [str(x).upper() for x in ncbi_df["assembly_acc"]] + # preserve the original column, but index ncbi_df on it as well ncbi_df = ncbi_df.set_index(pd.Index(list(ncbi_df[column]))) ## CHANGE TO ACCOMODATE BIOSAMPLE/OTHER NCBICOLUMNS if ncbi_column.lower() != "assembly": - vprint(f"{spacer}Assembling NCBI ftp directories", v=verbose, flush=True) + logger.debug(f"{spacer}Assembling NCBI ftp directories") acc2log = compile_log(output_path + "ncbiDwnld.log") acc2log, failed, ncbi_df = collect_assembly_accs( ncbi_df, @@ -439,14 +666,14 @@ def main( ) column = "assembly_acc" + # collect_assembly_accs supplies the column here, so it needs the same + # normalization before this index is taken from it + ncbi_df["assembly_acc"] = [str(x).upper() for x in ncbi_df["assembly_acc"]] ncbi_df = ncbi_df.set_index(pd.Index(list(ncbi_df[column]))) new_df = pd.DataFrame() ## GUARANTEE ASSEMBLY ACCESSIONS ARE LABELED THIS COLUMN NAME acc_file = output_path + "assembly_accs.txt" - ncbi_df["assembly_acc"] = list([x.upper() for x in ncbi_df["assembly_acc"]]) - with open(acc_file, "w") as out: - out.write("\n".join([str(x) for x in list(ncbi_df["assembly_acc"])])) include = "" req_files = set() @@ -470,91 +697,127 @@ def main( else: annotated = False - # Run downloads - count = 0 - while count < 3: - if not count: - vprint(f"{spacer}Downloading data", v=verbose, flush=True) - count += 1 - else: - count += 1 - vprint(f"{spacer}\tAttempt {count}", v=verbose, flush=True) - - run_datasets( - include, + # Chunk the accessions so datasets is called on `chunk` accessions at a time + all_accs = [str(x) for x in list(ncbi_df["assembly_acc"])] + acc_chunks = [all_accs[i : i + chunk] for i in range(0, len(all_accs), chunk)] + + # Run downloads chunk-by-chunk, accumulating results + acc2files, acc2org, failed = {}, {}, [] + dead_chunks, consecutive_dead = [], 0 + for chunk_i, acc_chunk in enumerate(acc_chunks): + if len(acc_chunks) > 1: + logger.debug( + f"{spacer}Chunk {chunk_i + 1}/{len(acc_chunks)} " + + f"({len(acc_chunk)} accessions)" + ) + c_acc2files, c_acc2org, c_failed = download_datasets( + acc_chunk, acc_file, + include, + req_files, + annotated, output_path, api=api, verbose=verbose, - annotated=annotated, - ) - - # Parse download output, add to df - acc2files, acc2org, failed = parse_datasets( - output_path + "ncbi_dataset.zip", output_path, req_files, spacer + spacer=spacer, ) - if acc2files == False and acc2org == False and failed == False: + if c_acc2files is False: + # one chunk NCBI will not serve is not worth discarding the chunks + # that succeeded; its accessions fall out of the set difference below + # and are reported as failures. Several in a row is the repository or + # the connection being gone, which the remaining chunks cannot fix + dead_chunks.append(chunk_i + 1) + consecutive_dead += 1 + if consecutive_dead >= 3: + logger.error( + f"{spacer}ncbiDwnld failed {consecutive_dead} consecutive chunks" + ) + sys.exit(10) + logger.warning( + f"{spacer}chunk {chunk_i + 1}/{len(acc_chunks)} failed; continuing" + ) continue - else: - break - - if acc2files == False and acc2org == False and failed == False: - eprint(f"{spacer}ERROR: ncbiDwnld failed {count} attempts", flush=True) - # maybe add a fallback to the old methodology here - eprint(f"{spacer}Consider --fallback", flush=True) - sys.exit(10) + consecutive_dead = 0 + acc2files.update(c_acc2files) + acc2org.update(c_acc2org) + failed.extend(c_failed) + + if dead_chunks: + logger.warning( + f"{spacer}{len(dead_chunks)}/{len(acc_chunks)} chunk(s) failed to " + + "download; their accessions are reported as failures" + ) + logger.debug(f"{spacer}failed chunks: {dead_chunks}") failed.extend( sorted(set(ncbi_df["assembly_acc"]).difference(set(acc2files.keys()))) ) - # Attempt RefSeq accessions + # Attempt the counterpart repository for whatever this one did not serve + acc2pair = {} if failed: - vprint( - f"{spacer}Attempting alternative repository for failed downloads", - v=verbose, - flush=True, - ) - reattempt_acc = [] - for acc in failed: - if acc.upper().startswith("GCA"): - reattempt_acc.append(acc.upper().replace("GCA_", "GCF_")) - elif acc.upper().startswith("GCF"): - reattempt_acc.append(acc.upper().replace("GCF_", "GCA_")) + logger.debug(f"{spacer}Attempting alternative repository for failed downloads") acc_file_re = output_path + "assembly_accs.reattempt.txt" - with open(acc_file_re, "w") as out: - out.write("\n".join(reattempt_acc)) - - run_datasets( - include, acc_file_re, output_path, verbose=verbose, annotated=annotated - ) - acc2files_r, acc2org_r, failed_r = parse_datasets( - output_path + "ncbi_dataset.zip", output_path, req_files + acc2pair = resolve_paired_accs( + failed, acc_file_re, output_path, api=api, spacer=spacer ) - acc2files = {**acc2files, **acc2files_r} - acc2org = {**acc2org, **acc2org_r} + pair2acc = {v: k for k, v in acc2pair.items()} + reattempt_acc = sorted(pair2acc) + unpaired = len(failed) - len(reattempt_acc) + if unpaired: + logger.debug( + f"{spacer}\t{unpaired} failed accession(s) are not mirrored in " + + "the alternative repository" + ) - failed = [] - for acc in failed_r: - if acc.upper().startswith("GCA"): - failed.append(acc.upper().replace("GCA_", "GCF_")) - elif acc.upper().startswith("GCF"): - failed.append(acc.upper().replace("GCF_", "GCA_")) + # Chunk the reattempt accessions as well + reattempt_chunks = [ + reattempt_acc[i : i + chunk] + for i in range(0, len(reattempt_acc), chunk) + ] + recovered = set() + for acc_chunk in reattempt_chunks: + c_acc2files, c_acc2org, c_failed = download_datasets( + acc_chunk, + acc_file_re, + include, + req_files, + annotated, + output_path, + api=api, + verbose=verbose, + spacer=spacer, + ) + if c_acc2files is False: + continue + acc2files = {**acc2files, **c_acc2files} + acc2org = {**acc2org, **c_acc2org} + recovered.update(pair2acc[x] for x in c_acc2files if x in pair2acc) + + # only what was actually retrieved leaves the failure list. A counterpart + # NCBI never served is absent from the archive rather than named in + # c_failed, so subtracting what came back is the only accounting that + # sees it -- rebuilding the list from c_failed instead dropped every + # unmirrored accession out of the run without a trace + failed = sorted(set(failed).difference(recovered)) # Parse download output, add to df # Report failed failed_set = set(failed) rep_failed = [] + + def report_failure(acc, row): + try: + rep_failed.append([acc, datetime.strftime(row["version"], "%Y%m%d")]) + except TypeError: + rep_failed.append([acc, row["version"]]) + for acc, row in ncbi_df.iterrows(): - if acc.startswith("GCA_"): - check_acc = acc.replace("GCA_", "GCF_") - elif acc.startswith("GCF_"): - check_acc = acc.replace("GCF_", "GCA_") + # the counterpart NCBI actually holds, which versions independently of + # acc; None when the assembly is not mirrored at all + check_acc = acc2pair.get(acc) if acc in failed_set: - try: - rep_failed.append([acc, datetime.strftime(row["version"], "%Y%m%d")]) - except TypeError: - rep_failed.append([acc, row["version"]]) + report_failure(acc, row) elif acc in acc2files: try: for file_t, file_p in acc2files[acc].items(): @@ -563,7 +826,7 @@ def main( ncbi_df.at[acc, tax] = name new_df = pd.concat([new_df, ncbi_df.loc[acc].to_frame().T]) except AttributeError: # multiple entries - eprint(f"{spacer}WARNING: {acc} is redundant", flush=True) + logger.warning(f"{spacer}{acc} is redundant") for acc1, row1 in ncbi_df.loc[acc].iterrows(): for file_t, file_p in acc2files[acc].items(): row1[file_t] = file_p @@ -579,12 +842,7 @@ def main( ncbi_df.at[acc, tax] = name new_df = pd.concat([new_df, ncbi_df.loc[acc].to_frame().T]) except AttributeError: # multiple entries - vprint( - f"{spacer}\tWARNING: {check_acc} is redundant", - v=verbose, - e=True, - flush=True, - ) + logger.debug(f"{spacer}\t{check_acc} is redundant") for acc1, row1 in ncbi_df.loc[acc].iterrows(): for file_t, file_p in acc2files[check_acc].items(): row1[file_t] = file_p @@ -592,6 +850,13 @@ def main( row1[tax] = name new_df = pd.concat([new_df, row1.to_frame().T]) + else: + # nothing retrieved it and nothing named it a failure. Without this + # the accession leaves the run in neither new_df nor rep_failed, + # which is how an initialization silently lost 214 of 525 genomes + logger.debug(f"{spacer}\t{acc} was neither retrieved nor reported") + report_failure(acc, row) + if "fna" in new_df.keys(): new_df = new_df.rename(columns={"fna": "assemblyPath"}) if "gff3" in new_df.keys(): @@ -600,7 +865,7 @@ def main( return new_df, rep_failed -def get_SRA(assembly_acc, fastqdump="fastq-dump", pe=True): +def get_sra(assembly_acc, fastqdump="fastq-dump", pe=True): handle = Entrez.esearch(db="SRA", term=assembly_acc) ids = Entrez.read(handle)["IdList"] @@ -609,7 +874,7 @@ def get_SRA(assembly_acc, fastqdump="fastq-dump", pe=True): records = Entrez.read(handle, validate=False) for record in records: srr = re.search(r'Run acc="(S\w+\d+)"', record["Runs"])[1] - print("\t\t" + srr, flush=True) + logger.info("\t\t" + srr) cmd, count = 1, 0 if pe: while cmd and count < 3: @@ -625,7 +890,7 @@ def get_SRA(assembly_acc, fastqdump="fastq-dump", pe=True): cmd = subprocess.call( [fastqdump, "--split-3", "--gzip", srr], stdout=subprocess.PIPE ) - if os.path.isfile(f"{srr}_1.fastq.gz"): + if Path(f"{srr}_1.fastq.gz").is_file(): # if os.path.isfile(srr + '_1.fastq'): # cmd = subprocess.call(['gzip', f'{srr}_1.fastq']) # cmd = subprocess.call(['gzip', f'{srr}_2.fastq']) @@ -637,9 +902,7 @@ def get_SRA(assembly_acc, fastqdump="fastq-dump", pe=True): ) else: # cmd = subprocess.call(['gzip', f'{srr}.fastq']) - print( - "\t\t\tWARNING: file failed or not paired-end", flush=True - ) + logger.warning("file failed or not paired-end") else: while cmd and count < 3: count += 1 @@ -655,25 +918,25 @@ def get_SRA(assembly_acc, fastqdump="fastq-dump", pe=True): if cmd: continue # cmd = subprocess.call(['gzip', f'{srr}.fastq']) - if os.path.isfile(f"{srr}.fastq.gz"): + if Path(f"{srr}.fastq.gz").is_file(): shutil.move(f"{srr}.fastq.gz", f"{assembly_acc}_{srr}.fq.gz") else: - print("\t\t\tERROR: file failed", flush=True) + logger.error("file failed") -def goSRA(df, output=os.getcwd() + "/", pe=True, column="sra"): +def go_sra(df, output=str(Path.cwd()) + "/", pe=True, column="sra"): print() sra_dir = output + "sra/" - if not os.path.isdir(sra_dir): - os.mkdir(sra_dir) + if not Path(sra_dir).is_dir(): + Path(sra_dir).mkdir() os.chdir(sra_dir) - fastqdump = findExecs("fastq-dump", exit={"fastq-dump"}) + fastqdump = find_execs("fastq-dump", exit={"fastq-dump"}) count = 0 for i, row in df.iterrows(): - print("\t" + row[column], flush=True) - get_SRA(row[column], fastqdump[0]) + logger.info("\t" + row[column]) + get_sra(row[column], fastqdump[0]) count += 1 if count >= 10: time.sleep(1) @@ -681,6 +944,9 @@ def goSRA(df, output=os.getcwd() + "/", pe=True, column="sra"): def cli(): + # BioPython (Bio.Entrez) raises a UserWarning when Entrez.email is unset; + # silence it so download output stays readable. + warnings.filterwarnings("ignore", category=UserWarning, module=r"Bio(\.|$)") parser = argparse.ArgumentParser( description="GenBank/RefSeq downloading utility. Downloads " + "accession by accession" @@ -713,70 +979,67 @@ def cli(): + "DEFAULT: attempt to decipher", ) parser.add_argument("-o", "--output", help="Output directory") - parser.add_argument("-e", "--email", help="NCBI email") parser.add_argument("--api", help="NCBI API key for high query rate") parser.add_argument( - "--fallback", action="store_true", help="Fallback mode if datasets fails" + "--chunk", + type=int, + default=25, + help="Accessions to download per datasets call; larger chunks are " + + "likelier to be reset mid-transfer by NCBI; DEFAULT: 25", ) args = parser.parse_args() + setup_logging(verbose=getattr(args, "verbose", False)) - if args.email: - ncbi_email = args.email - Entrez.email = ncbi_email - if args.api: - ncbi_api = args.api - Entrez.api_key = ncbi_api - else: - ncbi_api = None + if args.api: + ncbi_api = args.api else: - ncbi_email, ncbi_api, jgi_email, jgi_pwd = loginCheck(jgi=False) - Entrez.email = ncbi_email - if ncbi_api: - Entrez.api_key = ncbi_api + ncbi_api, jgi_email, jgi_pwd = login_check(jgi=False) + if ncbi_api: + Entrez.api_key = ncbi_api if not args.output: - output = mkOutput(None, "ncbiDwnld") + output = mk_output(None, "download_ncbi") else: output = format_path(args.output) - findExecs("datasets", exit={"datasets"}) + find_execs("datasets", exit={"datasets"}) args_dict = { "NCBI Table": args.input, - "email": ncbi_email, "Assemblies": args.assembly, "Proteomes": args.proteome, ".gff3's": args.gff3, "Transcripts": args.transcript, "SRA": args.sra, + "Chunk": args.chunk, } start_time = intro("Download NCBI files", args_dict) if args.sra: - if os.path.isfile(format_path(args.input)): + if Path(format_path(args.input)).is_file(): if not args.column: - goSRA( + go_sra( pd.read_csv(format_path(args.input), sep="\t", names=["sra"]), output, pe=args.paired, column="sra", ) else: - goSRA( + go_sra( pd.read_csv(format_path(args.input), sep="\t"), output, pe=args.paired, column=args.column, ) else: - goSRA( + go_sra( pd.DataFrame({"sra": split_input(args.input)}), output, pe=args.paired, column="sra", ) else: - if os.path.isfile(format_path(args.input)): + if Path(format_path(args.input)).is_file(): ncbi_df = pd.read_csv(args.input, sep="\t", header=None) if not args.column: if "assembly_acc" in ncbi_df.keys(): @@ -817,34 +1080,19 @@ def cli(): ncbi_column = "assembly" ncbi_df = ncbi_df.drop_duplicates(column) - if args.fallback: - from mycotools.ncbi_dwnld_fallback import main as main_fallback - - new_df, failed = main_fallback( - assembly=args.assembly, - column=column, - ncbi_column=ncbi_column, - proteome=args.proteome, - gff3=args.gff3, - transcript=args.transcript, - ncbi_df=ncbi_df, - output_path=output, - verbose=True, - spacer="", - ) - else: - new_df, failed = main( - assembly=args.assembly, - column=column, - ncbi_column=ncbi_column, - proteome=args.proteome, - gff3=args.gff3, - transcript=args.transcript, - ncbi_df=ncbi_df, - output_path=output, - verbose=True, - spacer="", - ) + new_df, failed = main( + assembly=args.assembly, + column=column, + ncbi_column=ncbi_column, + proteome=args.proteome, + gff3=args.gff3, + transcript=args.transcript, + ncbi_df=ncbi_df, + output_path=output, + verbose=True, + spacer="", + chunk=args.chunk, + ) new_df = new_df.rename(columns={"index": "#assembly_accession"}) new_df["source"] = "ncbi" new_df["useRestriction (yes/no)"] = "no" @@ -853,7 +1101,7 @@ def cli(): new_df.to_csv(output + "ncbiDwnld.predb", sep="\t", index=None) if failed: - eprint("ERROR: " + ",".join([str(x[0]) for x in failed]), flush=True) + logger.error(",".join([str(x[0]) for x in failed])) outro(start_time) diff --git a/mycotools/gff/__init__.py b/mycotools/gff/__init__.py new file mode 100644 index 0000000..9ac4936 --- /dev/null +++ b/mycotools/gff/__init__.py @@ -0,0 +1,36 @@ +#! /usr/bin/env python3 +"""Dispatcher for the `mycotools gff` subcommand. + +Routes `mycotools gff ...` to a gff manipulation/rendering module.""" +from mycotools.lib.subcmd import Dispatcher + +# subcommand name/alias -> submodule within this package (mycotools.gff.) +SUBCOMMANDS = { + "add": "add", + "svg": "svg", +} + +DESCRIPTION = """Manipulate and render gff3 files + +Tools (all following arguments are forwarded to the tool): + add add an entry to an existing gff + svg render a gff locus to SVG + +Examples: + mycotools gff add -h + mycotools gff svg -h""" + +_dispatcher = Dispatcher( + "mycotools gff", + "mycotools.gff", + SUBCOMMANDS, + DESCRIPTION, + metavar="TOOL", + arg_help="gff tool (see below)", +) +main = _dispatcher.main +cli = _dispatcher.cli + + +if __name__ == "__main__": + cli() diff --git a/mycotools/gff/__main__.py b/mycotools/gff/__main__.py new file mode 100644 index 0000000..0d8834e --- /dev/null +++ b/mycotools/gff/__main__.py @@ -0,0 +1,6 @@ +#! /usr/bin/env python3 +"""Enable ``python -m mycotools.gff`` to run the gff dispatcher.""" +from mycotools.gff import cli + +if __name__ == "__main__": + cli() diff --git a/mycotools/add2gff.py b/mycotools/gff/add.py similarity index 83% rename from mycotools/add2gff.py rename to mycotools/gff/add.py index ef00430..83a06b6 100755 --- a/mycotools/add2gff.py +++ b/mycotools/gff/add.py @@ -7,37 +7,40 @@ # NEED a protein_id import option -import os +import logging import re import sys import argparse from collections import defaultdict -from mycotools.lib.kontools import sys_start, format_path, eprint, mkOutput -from mycotools.lib.biotools import gff2list, list2gff, gff3Comps, gff2Comps, gtfComps -from mycotools.lib.dbtools import mtdb, primaryDB -from mycotools.utils.curGFF3 import rename_and_organize +from mycotools.lib.kontools import format_path, mk_output, setup_logging +from mycotools.lib.biotools import gff2list, list2gff, gff3_comps, gff2_comps, gtf_comps +from mycotools.lib.dbtools import mtdb, primary_db +from mycotools.utils.cur_gff3 import rename_and_organize +from pathlib import Path + +logger = logging.getLogger(__name__) def determine_version(toadd_gff, ome=None): """determine gff version for regex compilations""" for entry in toadd_gff: if entry["type"] == "gene": - if re.search(gff3Comps()["id"], entry["attributes"]): - return toadd_gff, gff3Comps() - elif re.search(gtfComps()["id"], entry["attributes"]): - return toadd_gff, gtfComps() - elif re.search(gff2Comps()["id"], entry["attributes"]): - return toadd_gff, gff2Comps() + if re.search(gff3_comps()["id"], entry["attributes"]): + return toadd_gff, gff3_comps() + elif re.search(gtf_comps()["id"], entry["attributes"]): + return toadd_gff, gtf_comps() + elif re.search(gff2_comps()["id"], entry["attributes"]): + return toadd_gff, gff2_comps() elif entry["type"] == "start_codon": # get this shit out - from mycotools.utils.gtf2gff3 import main as curAnn + from mycotools.utils.gtf2gff3 import main as cur_ann if not ome: - eprint("\nOme required for gtf input", flush=True) + logger.info("Ome required for gtf input") sys.exit(1) - new_gff = curAnn(toadd_gff, ome)[0] - return new_gff, gff3Comps() + new_gff = cur_ann(toadd_gff, ome)[0] + return new_gff, gff3_comps() else: # finished the for loop and no version detected, assume gff3 - return toadd_gff, gff3Comps() + return toadd_gff, gff3_comps() # eprint('\nCould not determine gff version', flush = True) @@ -50,7 +53,7 @@ def id_mtdb_accs(gff): mtdb_accs = [] for entry in gff: try: - mtdb_acc = re.search(gff3Comps()["Alias"], entry["attributes"])[1] + mtdb_acc = re.search(gff3_comps()["Alias"], entry["attributes"])[1] except TypeError: # no alias continue if "_manual" in mtdb_acc: # mtdb accession explicit @@ -100,7 +103,7 @@ def parse_toadd(toadd_gff, comps, ome, mtdb_acc=0): score, phase = gene["score"], gene["phase"] gene_acc = ome + "_manual" + str(mtdb_acc) if len(gene_dict["rna"]) > 1: - eprint("\nAlternately spliced loci currently not supported") + logger.info("Alternately spliced loci currently not supported") sys.exit(2) gene_dict["cds"] = sorted(gene_dict["cds"], key=lambda x: x[0]) @@ -197,7 +200,7 @@ def compile_mtdb_scaf(scafs, gff): for entry in gff: if entry["type"] == "gene": if entry["seqid"] in scafs: - mtdb_acc = re.search(gff3Comps()["Alias"], entry["attributes"])[1] + mtdb_acc = re.search(gff3_comps()["Alias"], entry["attributes"])[1] coord_tup = tuple(sorted([entry["start"], entry["end"]])) old_coords[entry["seqid"]][mtdb_acc] = sorted( [entry["start"], entry["end"]] @@ -218,7 +221,7 @@ def add_to_mtdb_gff(curadd_gff, addto_gff, ome, replace=False): if entry["type"] == "gene": start, stop = entry["start"], entry["end"] seqid = entry["seqid"] - new_acc = re.search(gff3Comps()["Alias"], entry["attributes"])[1] + new_acc = re.search(gff3_comps()["Alias"], entry["attributes"])[1] if seqid in ref_coords: for ref_coord, ref_acc in ref_coords[seqid].items(): ref_start, ref_stop = ref_coord @@ -229,12 +232,12 @@ def add_to_mtdb_gff(curadd_gff, addto_gff, ome, replace=False): ref_acc ] = new_acc # what about multiple gene # updates? e.g. fusions? this would delete - eprint(ref_acc + "\t->\t" + new_acc, flush=True) + logger.info(ref_acc + "->\t" + new_acc) out_gff = [] for entry in addto_gff: if entry["seqid"] in update: - mtdb_acc = re.search(gff3Comps()["Alias"], entry["attributes"])[1] + mtdb_acc = re.search(gff3_comps()["Alias"], entry["attributes"])[1] if mtdb_acc in update[entry["seqid"]]: continue out_gff.append(entry) @@ -259,19 +262,18 @@ def main(toadd_gff, addto_gff=[], ome=None, replace=False): def prep_mtdb_update(new_gff, ome, db): - from mycotools.predb2mtdb import main as predb2mtdb - from mycotools.lib.dbtools import mtdb, primaryDB + from mycotools.mtdb.predb import main as predb2mtdb - out_dir = mkOutput(format_path(os.getcwd()), "add2gff") + out_dir = mk_output(format_path(str(Path.cwd())), "gff_add") wrk_dir = out_dir + "working/" - if not os.path.isdir(wrk_dir): - os.mkdir(out_dir + "working/") - if not os.path.isdir(wrk_dir + "fna/"): - os.mkdir(out_dir + "working/fna/") - if not os.path.isdir(wrk_dir + "gff3/"): - os.mkdir(out_dir + "working/gff3/") - if not os.path.isdir(wrk_dir + "faa/"): - os.mkdir(out_dir + "working/faa/") + if not Path(wrk_dir).is_dir(): + Path(out_dir + "working/").mkdir() + if not Path(wrk_dir + "fna/").is_dir(): + Path(out_dir + "working/fna/").mkdir() + if not Path(wrk_dir + "gff3/").is_dir(): + Path(out_dir + "working/gff3/").mkdir() + if not Path(wrk_dir + "faa/").is_dir(): + Path(out_dir + "working/faa/").mkdir() new_gff_path = out_dir + "working/gff3/" + ome + ".new.gff3" with open(new_gff_path, "w") as out: @@ -298,7 +300,7 @@ def prep_mtdb_update(new_gff, ome, db): update_db_path = out_dir + "add2gff.mtdb" count = 1 - while os.path.isfile(update_db_path): + while Path(update_db_path).is_file(): update_db_path = re.sub(r"_\d+$", "", update_db_path) update_db_path += f"_{count}" count += 1 @@ -332,8 +334,9 @@ def cli(): action="store_true", help="[-a] Prepare output for mtdb update", ) - parser.add_argument("-d", "--mtdb", default=primaryDB()) + parser.add_argument("-d", "--mtdb", default=primary_db()) args = parser.parse_args() + setup_logging(verbose=getattr(args, "verbose", False)) # usage = 'Add gff to an existing mtdb gff.\n' \ # + 'add2gff.py ' @@ -342,10 +345,10 @@ def cli(): if not args.addto: if args.update: - eprint("\nERROR: -u requires -a", flush=True) + logger.error("-u requires -a") sys.exit(7) if not args.ome: - eprint("\nERROR: need -a or -o", flush=True) + logger.error("need -a or -o") sys.exit(5) else: ome = args.ome @@ -353,14 +356,14 @@ def cli(): else: ome = None addto_file = format_path(args.addto) - if not os.path.isfile(addto_file): - eprint("\nERROR: -a does not exist", flush=True) + if not Path(addto_file).is_file(): + logger.error("-a does not exist") sys.exit(6) addto_gff = gff2list(addto_file) toadd_file = format_path(args.input) - if not os.path.isfile(toadd_file): - eprint("\nERROR: -i does not exist", flush=True) + if not Path(toadd_file).is_file(): + logger.error("-i does not exist") sys.exit(7) toadd_gff = import_toadd_gff(toadd_file) diff --git a/mycotools/gff2svg.py b/mycotools/gff/svg.py similarity index 87% rename from mycotools/gff2svg.py rename to mycotools/gff/svg.py index 24db53e..e9d5add 100755 --- a/mycotools/gff2svg.py +++ b/mycotools/gff/svg.py @@ -3,17 +3,26 @@ # NEED TO PARSE FOR IN GENE COORDINATES AND ANNOTATIONS # NEED to create a single file output option for multiple inputs -import os import re import sys import random +import logging import argparse -from mycotools.lib.kontools import sys_start, format_path, file2list, getColors +from mycotools.lib.kontools import ( + format_path, + file2list, + get_colors, + setup_logging, +) from dna_features_viewer import GraphicFeature, GraphicRecord -from mycotools.lib.biotools import gff2list, gff3Comps +from mycotools.lib.biotools import gff2list, gff3_comps +from pathlib import Path -def compileProducts(gff, prod_comp, types={"tRNA", "mRNA", "rRNA"}): +logger = logging.getLogger(__name__) + + +def compile_products(gff, prod_comp, types={"tRNA", "mRNA", "rRNA"}): # find all the product attributes for color pallette selection products = [] for entry in gff: @@ -31,7 +40,7 @@ def gff2svg( svg_path, product_dict, colors, - prod_comp=gff3Comps()["product"], + prod_comp=gff3_comps()["product"], width=10, null="hypothetical protein", types={"tRNA", "mRNA", "rRNA"}, @@ -61,7 +70,7 @@ def gff2svg( elif not gen_new_colors: product = null elif product not in product_dict: - print(product) + logger.debug(product) try: # try to use the next color product_dict[product] = colors[count] count += 1 @@ -121,7 +130,7 @@ def main( svg_path, product_dict={}, width=10, - prod_comp=gff3Comps()["product"], + prod_comp=gff3_comps()["product"], null="hypothetical protein", types={"tRNA", "mRNA", "rRNA"}, labels=True, @@ -134,16 +143,16 @@ def main( if not wheel and not product_dict: if not set(product_dict.keys()).difference({null}): # if no keys or null is the only product key - products = compileProducts(gff_list, prod_comp, types=types) + products = compile_products(gff_list, prod_comp, types=types) else: products = list(product_dict.keys()) - colors = getColors(len(products)) + colors = get_colors(len(products)) elif wheel == 1: # spoof function to get wheel - colors = getColors(1) + colors = get_colors(1) elif wheel == 2: - colors = getColors(17) + colors = get_colors(17) elif wheel == 3: - colors = getColors(28) + colors = get_colors(28) else: colors = None @@ -207,9 +216,10 @@ def cli(): ) parser.add_argument("-o", "--output", help="Optional output directory") args = parser.parse_args() + setup_logging(verbose=getattr(args, "verbose", False)) if not args.regex: - regex = gff3Comps()["product"] + regex = gff3_comps()["product"] else: regex = r"" if args.regex.startswith(("'", '"')): @@ -227,16 +237,12 @@ def cli(): types = set(args.type.split()) if args.input: - if args.output: - out_dir = format_path(args.output) - else: - out_dir = format_path(os.path.dirname(args.input)) gffs = file2list(args.input) for entry in gffs: if entry.endswith("/"): entry = re.sub(r"/+$", "", entry) - svg_path = os.path.dirname(gffs[0]) + re.sub( - r"\.gf[^\.]+$", ".svg", os.path.basename(gffs[0]) + svg_path = str(Path(gffs[0]).parent) + re.sub( + r"\.gf[^\.]+$", ".svg", Path(gffs[0]).name ) product_dict = main( gff2list(gffs[0]), @@ -248,8 +254,8 @@ def cli(): shuffle=args.shuffle, ) for gff in gffs[1:]: - svg_path = os.path.dirname(gff) + re.sub( - r"\.gf[^\.]+$", ".svg", os.path.basename(gff) + svg_path = str(Path(gff).parent) + re.sub( + r"\.gf[^\.]+$", ".svg", Path(gff).name ) product_dict = main( gff2list(gff), @@ -267,11 +273,11 @@ def cli(): args.gff = re.sub(r"/+$", "", args.gff) if args.output: svg_path = format_path(args.output) + re.sub( - r"\.gf[^\.]+$", ".svg", os.path.basename(args.gff) + r"\.gf[^\.]+$", ".svg", Path(args.gff).name ) else: - svg_path = os.path.dirname(args.gff) + re.sub( - r"\.gf[^\.]+$", ".svg", os.path.basename(args.gff) + svg_path = str(Path(args.gff).parent) + re.sub( + r"\.gf[^\.]+$", ".svg", Path(args.gff).name ) main( gff2list(args.gff), diff --git a/mycotools/homology/__init__.py b/mycotools/homology/__init__.py new file mode 100644 index 0000000..f530b87 --- /dev/null +++ b/mycotools/homology/__init__.py @@ -0,0 +1,39 @@ +#! /usr/bin/env python3 +"""Dispatcher for the `mycotools homology` subcommand. + +Routes `mycotools homology ...` to a homology-search module: `db` +searches a query against the database (BLAST/mmseqs/hmmer), `fasta` runs +hmmsearch/nhmmer on a fasta and returns a fasta of hits.""" +from mycotools.lib.subcmd import Dispatcher + +# subcommand name/alias -> submodule within this package (mycotools.homology.) +SUBCOMMANDS = { + "db": "db", + "fasta": "fasta", + "fa": "fasta", +} + +DESCRIPTION = """Search query sequence(s) against the database or a fasta + +Methods (all following arguments are forwarded to the method): + db search a query against the database (BLAST/mmseqs/hmmer) + fasta (fa) hmmsearch/nhmmer a fasta and return a fasta of hits + +Examples: + mycotools homology db -h + mycotools homology fasta -h""" + +_dispatcher = Dispatcher( + "mycotools homology", + "mycotools.homology", + SUBCOMMANDS, + DESCRIPTION, + metavar="METHOD", + arg_help="homology-search method (see below)", +) +main = _dispatcher.main +cli = _dispatcher.cli + + +if __name__ == "__main__": + cli() diff --git a/mycotools/homology/__main__.py b/mycotools/homology/__main__.py new file mode 100644 index 0000000..0ba7476 --- /dev/null +++ b/mycotools/homology/__main__.py @@ -0,0 +1,6 @@ +#! /usr/bin/env python3 +"""Enable ``python -m mycotools.homology`` to run the homology dispatcher.""" +from mycotools.homology import cli + +if __name__ == "__main__": + cli() diff --git a/mycotools/db2search.py b/mycotools/homology/db.py similarity index 84% rename from mycotools/db2search.py rename to mycotools/homology/db.py index aa359d9..1f159e4 100755 --- a/mycotools/db2search.py +++ b/mycotools/homology/db.py @@ -12,38 +12,37 @@ # NEED nhmmer option # NEED to .tmp and move files -import os import re import sys -import copy -import datetime +import logging import argparse import subprocess import multiprocessing as mp -from io import StringIO from collections import defaultdict -from mycotools.db2files import soft_main as db2files from mycotools.lib.kontools import ( intro, outro, collect_files, multisub, - findExecs, + find_execs, untardir, - eprint, format_path, - mkOutput, + mk_output, tardir, inject_args, stdin2str, + setup_logging, ) -from mycotools.lib.dbtools import primaryDB, mtdb +from mycotools.lib.dbtools import primary_db, mtdb from mycotools.lib.biotools import dict2fa, fa2dict, fa2dict_str -# from mycotools.extractHmmsearch import main as exHmm -from mycotools.acc2fa import dbmain as acc2fa_db, famain as acc2fa_fa -from mycotools.utils.extractHmmsearch import main as exHmm -from mycotools.utils.extractHmmAcc import grabAccs, main as absHmm +# from mycotools.extract_hmmsearch import main as ex_hmm +from mycotools.mtdb.acc2.fa import dbmain as acc2fa_db, famain as acc2fa_fa +from mycotools.utils.extract_hmmsearch import main as ex_hmm +from mycotools.utils.extract_hmm_acc import grab_accs +from pathlib import Path + +logger = logging.getLogger(__name__) def compile_hmm_cmd(db, hmm_path, output, ome_set=set(), cpu=1): @@ -78,7 +77,7 @@ def compile_hmm_cmd(db, hmm_path, output, ome_set=set(), cpu=1): return cmd_tuples -def compileextractHmmCmd(db, args, output): +def compile_extract_hmm_cmd(db, args, output): """ Inputs: mycotools db, argparse arguments, and output path Outputs: tuples of arguments for `run_ex_hmm` @@ -100,16 +99,16 @@ def compileextractHmmCmd(db, args, output): def run_ex_hmm(args, hmmsearch_out, output): - ome = os.path.basename(hmmsearch_out).replace(".out", "") + ome = Path(hmmsearch_out).name.replace(".out", "") try: with open(hmmsearch_out, "r") as raw: data = raw.read() except FileNotFoundError: - eprint("\tWARNING: " + ome + " failed", flush=True) + logger.warning("\t" + ome + " failed") return ome, False if len(data) > 100: # check for data # check for data - hmm_data = exHmm( + hmm_data = ex_hmm( data, args[0], args[1], args[2], args[3], args[4], header=False ) out_dict = {} @@ -129,14 +128,13 @@ def run_ex_hmm(args, hmmsearch_out, output): # does this overwrite other hits? return ome, out_dict else: - eprint("\tWARNING: " + ome + " empty results", flush=True) + logger.warning("\t" + ome + " empty results") return ome, False def comp_hmm_acc2fa(db, q_dict, coords=True): cmd_tuples = [] - fa_dict = {ome: row["faa"] for ome, row in db.items()} for q in q_dict: cmd_tuples.append( ( @@ -175,9 +173,9 @@ def compile_mafft_cmds(output, faa_dir): fas = collect_files(faa_dir, "faa") # grab completed fastas cmds = [] for fa in fas: - acc = os.path.basename(fa)[:-4] + acc = Path(fa).name[:-4] align = output + "/aligns/" + acc + ".mafft.fasta" - if os.path.isfile(align): # check for data + if Path(align).is_file(): # check for data with open(align, "r") as raw: data = raw.read() if len(data) > 10: @@ -195,8 +193,7 @@ def compile_hmmalign_cmds(output, accessions): hmm = output + "/hmms/" + acc + ".hmm" align = output + "/aligns/" + acc + ".stockholm" conv = output + "/aligns/" + acc + ".phylip" - trim = output + "/trimmed/" + acc + ".clipkit.fa" - if os.path.isfile(conv): + if Path(conv).is_file(): with open(conv, "r") as raw: data = raw.read() if len(data) > 10: @@ -233,17 +230,10 @@ def compile_trim_cmd(output, mod="", trimmed=None, ex="phylip"): cmd_tuples = [] aligns = collect_files(output + "aligns/", ex) if trimmed: - trimmed = set(os.path.basename(x).replace(".clipkit", "") for x in trimmed) - aligns = [ - x for x in aligns if os.path.basename(x).replace(ex, "") not in trimmed - ] + trimmed = set(Path(x).name.replace(".clipkit", "") for x in trimmed) + aligns = [x for x in aligns if Path(x).name.replace(ex, "") not in trimmed] for align in aligns: - trim = ( - f"{output}trimmed/" - + os.path.basename(align).replace(ex, "clipkit") - + "." - + ex - ) + trim = f"{output}trimmed/" + Path(align).name.replace(ex, "clipkit") + "." + ex args = ["clipkit", align, "-o", trim] if mod_args[0]: args.extend(mod_args) @@ -257,7 +247,7 @@ def compile_hmm_queries(hmm_paths, hmm_out): for hmm_path in hmm_paths: with open(hmm_path, "r") as hmm_raw: hmm_data = hmm_raw.read() - queries.extend(grabAccs(hmm_data)) + queries.extend(grab_accs(hmm_data)) complete_hmm += hmm_data + "\n" with open(hmm_out, "w") as hmm_oh: hmm_oh.write(complete_hmm.rstrip()) @@ -288,28 +278,28 @@ def hmmer_main( queries = compile_hmm_queries(hmm_paths, hmm_out) ome_set, skip1 = set(), False - if os.path.isdir(faa_dir): # is there a previous run? - print("\nCompiling previous run", flush=True) + if Path(faa_dir).is_dir(): # is there a previous run? + logger.info("Compiling previous run") # check if all fastas are made fas = collect_files(faa_dir, "faa") # grab completed fastas - ranQueries = [os.path.basename(fa).replace(".faa", "") for fa in fas] + ranQueries = [Path(fa).name.replace(".faa", "") for fa in fas] if not set(queries).difference(set(ranQueries)): skip1 = True if not skip1: # do not skip the first step - if os.path.isfile(output + "omes.tar.gz"): - if not os.path.isdir(ome_dir): + if Path(output + "omes.tar.gz").is_file(): + if not Path(ome_dir).is_dir(): untardir(output + "omes.tar.gz") # check what reports have been generated omes = collect_files(ome_dir, "out") - ome_set = set(os.path.basename(x).replace(".out", "") for x in omes) + ome_set = set(Path(x).name.replace(".out", "") for x in omes) else: - print("\thmmsearch -> hits.faa DONE", flush=True) + logger.info("\thmmsearch -> hits.faa DONE") if not skip1: - print("\nRunning " + os.path.basename(binary), flush=True) - if not os.path.isdir(ome_dir): - os.mkdir(ome_dir) + logger.info("Running " + Path(binary).name) + if not Path(ome_dir).is_dir(): + Path(ome_dir).mkdir() # run hmmer par_runs = round(((cpu - 1) / 2) - 0.5) @@ -326,17 +316,17 @@ def hmmer_main( ) for i, code in enumerate(hmmsearch_codes): if code: - eprint("\tERROR: " + str(hmmsearch_tuples[i]), flush=True) + logger.error("\t" + str(hmmsearch_tuples[i])) # extract results - print("\nExtracting hmmsearch output", flush=True) + logger.info("Extracting hmmsearch output") exHmm_args = [accessions, max_hits, query_cov, evalue, bitscore] - exHmm_tuples = compileextractHmmCmd(db, exHmm_args, ome_dir) + exHmm_tuples = compile_extract_hmm_cmd(db, exHmm_args, ome_dir) with mp.get_context("spawn").Pool(processes=cpu) as pool: hmmAligns = pool.starmap(run_ex_hmm, exHmm_tuples) mp.Process(target=tardir, args=[ome_dir]) - print("\nCompiling fastas", flush=True) + logger.info("Compiling fastas") # q_dict = {query: ome: alignment} q_dict = defaultdict(dict) for ome, hits in hmmAligns: @@ -353,7 +343,7 @@ def hmmer_main( return fa_dicts -def compileBlastCmd(ome, biofile, out_dir, blast_scaf): +def compile_blast_cmd(ome, biofile, out_dir, blast_scaf): return blast_scaf + ["-out", out_dir + ome + ".tsv", "-subject", biofile] @@ -390,12 +380,14 @@ def comp_blast_tups( for i, ome in enumerate(seq_db["ome"]): if seq_db[biotype][i]: blast_cmds.append( - " ".join(compileBlastCmd(ome, seq_db[biotype][i], out_dir, blast_scaf)) + " ".join( + compile_blast_cmd(ome, seq_db[biotype][i], out_dir, blast_scaf) + ) ) return blast_cmds -def compileDiamondCmd(ome, dmnd_db, out_dir, blast_scaf): +def compile_diamond_cmd(ome, dmnd_db, out_dir, blast_scaf): return blast_scaf + ["--out", out_dir + ome + ".tsv", "--subject", dmnd_db] @@ -414,8 +406,8 @@ def comp_diamond_tups( search_args=[], ): - if not os.path.isdir(out_dir + "dmnd/"): - os.mkdir(out_dir + "dmnd/") + if not Path(out_dir + "dmnd/").is_dir(): + Path(out_dir + "dmnd/").mkdir() blast_scaf = [ diamond, blast_type, @@ -461,7 +453,7 @@ def comp_diamond_tups( ] ) blast_cmds.append( - compileDiamondCmd(ome, out_dir + "dmnd/" + ome, out_dir, blast_scaf) + compile_diamond_cmd(ome, out_dir + "dmnd/" + ome, out_dir, blast_scaf) ) return db_cmds, blast_cmds @@ -480,8 +472,8 @@ def run_mmseq( ): db_dir = format_path("$MYCOGFF3/../db") - if not os.path.isdir(f"{out_dir}db/"): - os.mkdir(f"{out_dir}db/") + if not Path(f"{out_dir}db/").is_dir(): + Path(f"{out_dir}db/").mkdir() db_dir = format_path("$MYCOGFF3/../db/") createdb_cmds = [] db_path = f"{out_dir}db/searchdb" @@ -490,7 +482,7 @@ def run_mmseq( for i, ome in enumerate(seq_db["ome"]): db_path = f"{db_dir}{ome}_{biotype}" out_file = out_dir + ome + ".tsv" - if not os.path.isfile(db_path + ".dbtype"): + if not Path(db_path + ".dbtype").is_file(): # will fail at fastas that dont have sequences on one line createdb_cmds.append( ( @@ -506,8 +498,8 @@ def run_mmseq( ) if createdb_cmds: - print(f"\nCreating {len(createdb_cmds)} mmseqs search dbs", flush=True) - createdb_outs = multisub(createdb_cmds, processes=cpus, verbose=2) + logger.info(f"Creating {len(createdb_cmds)} mmseqs search dbs") + multisub(createdb_cmds, processes=cpus, verbose=2) # if len(query) > 1: # if not os.path.isfile(f'{out_dir}db/query.dbtype'): @@ -520,20 +512,20 @@ def run_mmseq( # query = [f'{out_dir}db/query'] # need to adjust check # create a concatenated mmseqs db for the search target - if not os.path.isfile(f"{out_dir}db/searchdb.dbtype"): + if not Path(f"{out_dir}db/searchdb.dbtype").is_file(): mergedbs_cmd = ["mmseqs", "mergedbs"] mergedbs_cmd.extend([f"{db_dir}{ome}_{biotype}" for ome in seq_db["ome"]]) mergedbs_cmd.insert(3, f"{out_dir}db/searchdb") - print("\nMerging search dbs", flush=True) - mergedbs_out = subprocess.call(mergedbs_cmd) # , stderr = subprocess.DEVNULL, + logger.info("Merging search dbs") + subprocess.call(mergedbs_cmd) # , stderr = subprocess.DEVNULL, # stdout = subprocess.DEVNULL) - print("\nSearching", flush=True) + logger.info("Searching") for i, q in enumerate(query): out_file = f"{out_dir}{q}.tsv" - if os.path.isfile(out_file): + if Path(out_file).is_file(): continue - print("\t" + q, flush=True) + logger.info("\t" + q) search_cmd = [ mmseqs, "search", @@ -552,7 +544,7 @@ def run_mmseq( if coverage: search_cmd.extend(["-c", str(coverage)]) - search_out = subprocess.call(search_cmd) # , stderr = subprocess.DEVNULL, + subprocess.call(search_cmd) # , stderr = subprocess.DEVNULL, # stdout = subprocess.DEVNULL) results_cmd = [ mmseqs, @@ -564,12 +556,12 @@ def run_mmseq( "--format-output", "qset,target,pident,tstart,tend,evalue,bits", ] - results_out = subprocess.call( + subprocess.call( results_cmd, stderr=subprocess.DEVNULL, stdout=subprocess.DEVNULL ) -def parseOutput( +def parse_output( algorithm, ome, file_, @@ -582,7 +574,7 @@ def parseOutput( ): ome_results = [ome, []] - if os.path.exists(file_): + if Path(file_).exists(): with open(file_, "r") as raw: data = [x.rstrip().split("\t") for x in raw if x.rstrip()] byq = defaultdict(list) @@ -604,12 +596,12 @@ def parseOutput( return ome_results -def parseOutput_mmseqs( +def parse_output_mmseqs( algorithm, ome, file_, bitscore=0, pident=0, evalue=0, max_hits=None, ppos=None ): ome_results = [ome, []] - if os.path.exists(file_): + if Path(file_).exists(): with open(file_, "r") as raw: data = [x.rstrip().split("\t") for x in raw if x.rstrip()] byq = defaultdict(list) @@ -630,14 +622,13 @@ def parseOutput_mmseqs( return ome_results -def compileResults(res_dict, skip=[]): +def compile_results(res_dict, skip=[]): output_res = {} for i in res_dict: for hit in res_dict[i]: query = hit[0] subject = hit[1] - pident = hit[2] start = hit[-4] end = hit[-3] if query not in output_res: @@ -656,7 +647,6 @@ def compileResults(res_dict, skip=[]): def comp_mmseq_acc2fa(db, biotype, output_res, coords=False, skip=None): cmd_tuples = [] - fa_dict = {ome: row["faa"] for ome, row in db.set_index().items()} for q, ome_dict in output_res.items(): q_accs = [] if coords: @@ -713,39 +703,39 @@ def comp_blast_acc2fa(db, biotype, output_res, coords=False, skip=None): try: acc2fa_cmds[query].append([list(set(accs)), db[i][biotype]]) except KeyError: - eprint("\t" + i + " not in db") + logger.warning("\t" + i + " not in db") return acc2fa_cmds -def prepOutput(out_dir): +def prep_output(out_dir): out_dir = format_path(out_dir) if not out_dir.endswith("/"): out_dir += "/" - if not os.path.isdir(out_dir): - os.mkdir(out_dir) + if not Path(out_dir).is_dir(): + Path(out_dir).mkdir() report_dir = out_dir + "reports/" - if not os.path.isdir(report_dir): - os.mkdir(report_dir) + if not Path(report_dir).is_dir(): + Path(report_dir).mkdir() return report_dir def run_denovo(report_dir, log_list0, log_name): - if os.path.isdir(report_dir): + if Path(report_dir).is_dir(): count = 0 - while os.path.isdir(report_dir[:-1] + str(count)): + while Path(report_dir[:-1] + str(count)).is_dir(): count += 1 report_dir = report_dir[:-1] + str(count) + "/" log_list0[0] = report_dir - os.mkdir(report_dir) + Path(report_dir).mkdir() with open(log_name, "w") as out: out.write("\n".join(log_list0)) return log_list0, report_dir -def db2searchLog( +def db2search_log( report_dir, blast, query, max_hits, evalue, bit, pident, coverage, out_dir, ppos ): log_list0 = [ @@ -761,31 +751,31 @@ def db2searchLog( ] prev, reparse = False, False - log_name = out_dir + "." + os.path.basename(out_dir[:-1]) + ".log" + log_name = out_dir + "." + Path(out_dir[:-1]).name + ".log" log_list1 = None - if not os.path.isfile(log_name): # generate a new log + if not Path(log_name).is_file(): # generate a new log with open(log_name, "w") as out: out.write("\n".join(log_list0)) else: # check the old one with open(log_name, "r") as raw: log_list1 = [x.rstrip() for x in raw if x] if blast != log_list1[1]: - eprint("\tInconsistent search algorithm, rerunning", flush=True) + logger.warning("\tInconsistent search algorithm, rerunning") log_list0, report_dir = run_denovo(report_dir, log_list0, log_name) elif blast == "mmseqs": if log_list1[-1] != log_list0[-1]: # coverage is off, need a rerun - eprint("\tCoverage changed, rerunning", flush=True) + logger.warning("\tCoverage changed, rerunning") log_list0, report_dir = run_denovo(report_dir, log_list0, log_name) # need to reparse if anything is different elif any(log_list0[i] != log_list1[i] for i in range(len(log_list0))): - eprint("\tDeleting old report compilations", flush=True) + logger.info("\tDeleting old report compilations") reports = collect_files(report_dir, "tsv") for r in reports: - os.remove(r) + Path(r).unlink() reparse = True prev = True elif log_list1[1] != log_list0[1] and log_list1[3:] != log_list0[3:]: - eprint("\tInconsistent thresholds, rerunning", flush=True) + logger.warning("\tInconsistent thresholds, rerunning") log_list0, report_dir = run_denovo(report_dir, log_list0, log_name) else: prev = True @@ -800,14 +790,14 @@ def prepare_search_run( prev, finished, rundb = False, set(), db if isinstance(query, list): query = ",".join(query) - log_list0, log_list1, prev, reparse = db2searchLog( + log_list0, log_list1, prev, reparse = db2search_log( report_dir, blast, query, max_hits, evalue, bit, pident, coverage, out_dir, ppos ) report_dir = log_list0[0] if prev: # reparse = False reports = collect_files(report_dir, "tsv") - finished = {os.path.basename(x)[:-4] for x in reports if os.path.getsize(x) > 0} + finished = {Path(x).name[:-4] for x in reports if Path(x).stat().st_size > 0} rundb, checkdb = mtdb({}).set_index("ome"), db.set_index("ome") for ome, val in checkdb.items(): if ome not in finished: @@ -821,10 +811,10 @@ def prepare_search_run( def prep_mmseq_output(rundb, report_dir, queries, convert=False): ome_res = defaultdict(str) for i, q in enumerate(queries): - if not os.path.isfile(f"{report_dir}{q}.tsv"): + if not Path(f"{report_dir}{q}.tsv").is_file(): continue with open(f"{report_dir}{q}.tsv", "r") as raw: - base = os.path.basename(q) + base = Path(q).name if not convert: for line in raw: line_d = line.split() @@ -845,7 +835,7 @@ def prep_mmseq_output(rundb, report_dir, queries, convert=False): def comp_mmseq_res(rundb, report_dir, queries, convert=False): for ome in rundb["ome"]: out_file = report_dir + ome + ".tsv" - if os.path.isfile(out_file): + if Path(out_file).is_file(): continue out_str, todel = "", [] for i, q in enumerate(queries): @@ -853,7 +843,7 @@ def comp_mmseq_res(rundb, report_dir, queries, convert=False): if not convert: out_str += raw.read().rstrip() + "\n" else: - base = os.path.basename(q) + base = Path(q).name for line in raw: line_d = line.split() line_d[0] = base @@ -863,10 +853,10 @@ def comp_mmseq_res(rundb, report_dir, queries, convert=False): out.write(out_str.rstrip()) for todel_file in todel: - os.remove(todel_file) + Path(todel_file).unlink() -def ObyOsearch( +def o_by_o_search( db, rundb, blast, @@ -887,7 +877,7 @@ def ObyOsearch( ppos=0, ): if len(rundb) > 0: - print("\nSearching on an ome-by-ome basis", flush=True) + logger.info("Searching on an ome-by-ome basis") if diamond: db_tups, search_tups = comp_diamond_tups( rundb, @@ -902,9 +892,9 @@ def ObyOsearch( coverage=coverage * 100, search_args=search_arg, ) - db_outs = multisub(db_tups, processes=cpus) - print(f"\t{len(search_tups)} searches to run", flush=True) - search_outs = multisub( + multisub(db_tups, processes=cpus) + logger.info(f"\t{len(search_tups)} searches to run") + multisub( search_tups, processes=cpus, verbose=2, injectable=True ) scale = 100000 @@ -921,7 +911,7 @@ def ObyOsearch( coverage=coverage * 100, search_args=search_arg, ) - search_outs = multisub( + multisub( search_tups, processes=cpus, verbose=2, shell=True, injectable=True ) scale = 100000 @@ -944,9 +934,9 @@ def ObyOsearch( ] ) - print("\nParsing reports", flush=True) + logger.info("Parsing reports") with mp.get_context("spawn").Pool(processes=cpus) as pool: - results = pool.starmap(parseOutput, tuple(parse_tups)) + results = pool.starmap(parse_output, tuple(parse_tups)) results_dict = {x[0]: x[1] for x in results} return results_dict @@ -981,10 +971,10 @@ def mmseqs_mngr( cpus=cpus, ) - print("\nExtracting ome reports", flush=True) + logger.info("Extracting ome reports") prep_mmseq_output(rundb, report_dir, query, convert=convert) - print("\nParsing output", flush=True) + logger.info("Parsing output") # prepare report parsing commands for multiprocessing parse_tups = [] for i, ome in enumerate(db["ome"]): @@ -1002,32 +992,32 @@ def mmseqs_mngr( ) with mp.get_context("spawn").Pool(processes=cpus) as pool: - results = pool.starmap(parseOutput_mmseqs, tuple(parse_tups)) + results = pool.starmap(parse_output_mmseqs, tuple(parse_tups)) results_dict = {x[0]: x[1] for x in results} return results_dict -def checkSearchDB(binary="blast"): +def check_search_db(binary="blast"): - db_date = os.path.basename(primaryDB()) + db_date = Path(primary_db()).name if "blast" in binary: search_db = format_path("$MYCOFAA/blastdb/" + db_date + ".00.psd") - if os.path.isfile(search_db): + if Path(search_db).is_file(): return search_db[:-7] - elif os.path.isfile(search_db[:-7] + ".psd"): + elif Path(search_db[:-7] + ".psd").is_file(): return search_db[:-7] else: search_db = format_path( "$MYCOFAA/blastdb/" + db_date.replace(".db", "") + ".mmseqs.db" ) - if os.path.isfile(search_db): + if Path(search_db).is_file(): return search_db def db_blast(db_path, blast_type, query, evalue, hsps, cpus, report_dir, diamond=False): - out_file = report_dir + os.path.basename(db_path)[:-3] + ".out" + out_file = report_dir + Path(db_path).name[:-3] + ".out" if not diamond: blast_scaf = [ blast_type, @@ -1084,9 +1074,7 @@ def db_blast(db_path, blast_type, query, evalue, hsps, cpus, report_dir, diamond def dbmmseq(db_path, query, evalue, cpus, report_dir, mmseqs="mmseqs", mem=None): - out_file = ( - report_dir + os.path.basename(db_path)[:-3].replace(".mmseqs", "") + ".out" - ) + out_file = report_dir + Path(db_path).name[:-3].replace(".mmseqs", "") + ".out" # output_str = '"query,target,pident,alen,mismatch,gapopen,qstart,qend,sstart,send,evalue,bits"' cmd_scaf = [ mmseqs, @@ -1113,7 +1101,7 @@ def dbmmseq(db_path, query, evalue, cpus, report_dir, mmseqs="mmseqs", mem=None) return cmd_call, out_file -def parseDBout(db, file_, bitscore=0, pident=0, ppos=0, max_hits=None): +def parse_db_out(db, file_, bitscore=0, pident=0, ppos=0, max_hits=None): ome_results = {} with open(file_, "r") as raw: @@ -1129,8 +1117,6 @@ def parseDBout(db, file_, bitscore=0, pident=0, ppos=0, max_hits=None): ome_results[ome] = [] ome_results[ome].append(data) - x_omes = set(db["ome"]) - omes_results = {x: ome_results[x] for x in ome_results if x in x_omes} if max_hits: out_results = {} @@ -1164,7 +1150,7 @@ def mmseqs_main( iterations=3, ): - report_dir = prepOutput(out_dir) + report_dir = prep_output(out_dir) if isinstance(query, str): query = [query] elif isinstance(query, dict): @@ -1202,8 +1188,8 @@ def mmseqs_main( reparse=reparse, ) - print("\nCompiling fastas", flush=True) - output_res = compileResults(results_dict, skip) + logger.info("Compiling fastas") + output_res = compile_results(results_dict, skip) output_fas = {} acc2fa_cmds = comp_mmseq_acc2fa( db, biotype, output_res, coords=coordinate, skip=None @@ -1237,18 +1223,16 @@ def blast_main( ): if blast in {"tblastn", "blastp"}: - seq_type = "prot" biotype = "faa" elif blast in {"blastx", "blastn"}: - seq_type = "nucl" biotype = "fna" else: - eprint("\nERROR: invalid search binary: " + blast, flush=True) + logger.error("invalid search binary: " + blast) # if blastdb: # insert function to make blastdb - report_dir = prepOutput(out_dir) + report_dir = prep_output(out_dir) if isinstance(query, str): query = [query] elif isinstance(query, dict): @@ -1264,14 +1248,14 @@ def blast_main( query = out_dir + "query.fa" if blastdb and not force: - print("\nSearching using MycotoolsDB searchdb", flush=True) + logger.info("Searching using MycotoolsDB searchdb") search_exit, search_output = db_blast( blastdb, blast, query, evalue, hsps, cpus, report_dir, diamond=diamond ) if search_exit: - eprint("\nERROR: search failed: " + str(search_exit)) + logger.error("search failed: " + str(search_exit)) sys.exit(10) - results_dict = parseDBout( + results_dict = parse_db_out( db, search_output, bitscore=bitscore, @@ -1293,7 +1277,7 @@ def blast_main( coverage, ppos, ) - results_dict = ObyOsearch( + results_dict = o_by_o_search( db, rundb, blast, @@ -1314,16 +1298,15 @@ def blast_main( ppos=ppos, ) - print("\nCompiling fastas", flush=True) - output_res = compileResults(results_dict, skip) + logger.info("Compiling fastas") + output_res = compile_results(results_dict, skip) output_fas = {} acc2fa_cmds = comp_blast_acc2fa( db, biotype, output_res, coords=coordinate, skip=None ) - queryfa = fa2dict(query) for query1, cmd in acc2fa_cmds.items(): output_fas[query1] = {} - print("\t" + query1, flush=True) + logger.info("\t" + query1) with mp.get_context("spawn").Pool(processes=cpus) as pool: results = pool.starmap(acc2fa_fa, acc2fa_cmds[query1]) for x in results: @@ -1348,7 +1331,7 @@ def cli(): ) i_arg = parser.add_argument_group("Inputs") - i_arg.add_argument("-d", "--mtdb", default=primaryDB()) + i_arg.add_argument("-d", "--mtdb", default=primary_db()) i_arg.add_argument( "-q", "--query", help='Profile database, sequence, or "-" for stdin fasta' ) @@ -1436,19 +1419,18 @@ def cli(): # parser.add_argument( '-c', '--coverage', type = float, help = 'Query coverage +/-, e.g. 0.5' ) # parser.add_argument('-f', '--force', action = 'store_true', help = 'Force ome-by-ome blast') args = parser.parse_args() + setup_logging(verbose=getattr(args, "verbose", False)) if args.algorithm not in algorithms: - eprint("\nERROR: {args.algorithm} not implemented", flush=True) + logger.error("{args.algorithm} not implemented") sys.exit(19) # query parsing if not args.query and not args.query_dir and not args.query_file: - eprint( - "\nERROR: --query, --query_dir, or --query_file not specified", flush=True - ) + logger.error("--query, --query_dir, or --query_file not specified") sys.exit(18) elif [args.query, args.query_dir, args.query_file].count(True) > 1: - eprint("\nERROR: multiple query types specified", flush=True) + logger.error("multiple query types specified") sys.exit(20) else: if args.query: @@ -1459,7 +1441,7 @@ def cli(): elif args.query_dir: queries = [ format_path(args.query_dir) + x - for x in os.listdir(format_path(args.query_dir)) + for x in [p.name for p in Path(format_path(args.query_dir)).iterdir()] ] else: with open(format_path(args.query_file), "r") as raw: @@ -1469,7 +1451,7 @@ def cli(): # diamond specific if args.diamond: if "blast" not in args.algorithm: - eprint(f"\nERROR: --diamond incompatible with {args.algorithm}", flush=True) + logger.error(f"--diamond incompatible with {args.algorithm}") sys.exit(4) del deps[0] deps.append("diamond") @@ -1477,7 +1459,7 @@ def cli(): # mmseqs-specific if args.algorithm == "mmseqs": if not args.seqtype or args.seqtype not in {"aa", "nt"}: - eprint("\nERROR: -st required for mmseqs", flush=True) + logger.error("-st required for mmseqs") sys.exit(3) if args.seqtype == "aa": biotype = "faa" @@ -1493,23 +1475,23 @@ def cli(): # queries = sorted(query_set) else: biotype = None - findExecs(deps, exit=set(deps)) + find_execs(deps, exit=set(deps)) # identity if args.identity: if args.algorithm == "hmmsearch": - eprint(f"\nERROR: --identity incompatible with hmmsearch", flush=True) + logger.error(f"--identity incompatible with hmmsearch") sys.exit(21) if not args.output: - base = os.getcwd() + "/" - output = mkOutput(base, "db2search") + base = str(Path.cwd()) + "/" + output = mk_output(base, "homology_db") else: base = format_path(args.output, force_dir=True) output = base - if not os.path.isdir(output): - os.mkdir(output) - # output = mkOutput(base, 'db2search') + if not Path(output).is_dir(): + Path(output).mkdir() + # output = mk_output(base, 'homology_db') if args.cpu and args.cpu < mp.cpu_count(): cpu = args.cpu @@ -1598,8 +1580,8 @@ def cli(): convert=args.convert, iterations=args.iterations, ) - if not os.path.isdir(output + "fastas/"): - os.mkdir(output + "fastas/") + if not Path(output + "fastas/").is_dir(): + Path(output + "fastas/").mkdir() for query in output_fas: with open(output + "fastas/" + query + ".search.fa", "w") as out: diff --git a/mycotools/fa2hmmer2fa.py b/mycotools/homology/fasta.py similarity index 68% rename from mycotools/fa2hmmer2fa.py rename to mycotools/homology/fasta.py index 6b294a6..d4b2098 100755 --- a/mycotools/fa2hmmer2fa.py +++ b/mycotools/homology/fasta.py @@ -2,6 +2,7 @@ # NEED to ditch extracthmm and move to simplified output parsing +import logging import os import re import sys @@ -9,16 +10,18 @@ import datetime import subprocess import multiprocessing as mp -from mycotools.extractHmmsearch import main as exHmm, grabNames -from mycotools.extractHmmAcc import main as extr_hmm -from mycotools.db2search import compAcc2fa -from mycotools.acc2fa import famain as acc2fa -from mycotools.lib.kontools import intro, outro, findExecs, eprint, format_path -from mycotools.lib.dbtools import mtdb, primaryDB +from mycotools.utils.extract_hmmsearch import main as ex_hmm, grab_names +from mycotools.utils.extract_hmm_acc import main as extr_hmm +from mycotools.homology.db import comp_hmm_acc2fa, hmm_acc2fa +from mycotools.lib.kontools import intro, outro, find_execs, format_path, setup_logging +from mycotools.lib.dbtools import mtdb, primary_db from mycotools.lib.biotools import dict2fa +from pathlib import Path +logger = logging.getLogger(__name__) -def runextractHmmAcc(hmm, accession, output): + +def run_extract_hmm_acc(hmm, accession, output): with open(hmm, "r") as raw: hmm_data = raw.read() @@ -29,7 +32,7 @@ def runextractHmmAcc(hmm, accession, output): return output -def runHmmer(fasta, hmm, output, cpu=1, binary="hmmsearch"): +def run_hmmer(fasta, hmm, output, cpu=1, binary="hmmsearch"): hmm_status = subprocess.call( [binary, "-o", output, "--cpu", str(cpu), hmm, fasta], @@ -44,11 +47,11 @@ def run_extract_hmm(hmm_out, top_hits, cov_threshold, evalue, query=True, acc=No with open(hmm_out, "r") as raw: data = raw.read() - accs = grabNames(data, query=query) + accs = grab_names(data, query=query) if len(accs) > 1: - hmm_data = exHmm(data, True, top_hits, cov_threshold, evalue, query=query) + hmm_data = ex_hmm(data, True, top_hits, cov_threshold, evalue, query=query) else: - hmm_data = exHmm( + hmm_data = ex_hmm( data, list(accs)[0], top_hits, cov_threshold, evalue, query=query ) @@ -77,18 +80,19 @@ def parse_hmm_data(hmm_data): def run_acc2fa(db, biotype, output_res, subhit=True, cpu=1): - acc2fa_cmds = compAcc2fa(db, biotype, output_res, subhit) - output_fas = {} - for query in acc2fa_cmds: - print("\t" + query, flush=True) - with mp.get_context("spawn").Pool(processes=cpu) as pool: - results = pool.starmap(acc2fa, acc2fa_cmds[query]) - output_fas[query] = "\n".join([dict2fa(x) for x in results]) + # comp_hmm_acc2fa builds one (db, {ome: [[seq, start, end], ...]}, query, + # coords) tuple per query; hmm_acc2fa turns each into (query, fa_dict) by + # retrieving from the proteome (faa). NOTE: this pair ignores `biotype`, so + # nhmmer (fna) hits are still pulled from faa -- use comp_blast_acc2fa if + # nucleotide retrieval is needed. + acc2fa_tuples = comp_hmm_acc2fa(db, output_res, coords=subhit) + with mp.get_context("spawn").Pool(processes=cpu) as pool: + fa_info = pool.starmap(hmm_acc2fa, acc2fa_tuples) - return output_fas + return {query: dict2fa(fa_dict) for query, fa_dict in fa_info} -def outputFas(output_fas, output_dir, fastaname): +def output_fas(output_fas, output_dir, fastaname): for query in output_fas: with open(output_dir + fastaname + "_" + query + ".fa", "w") as out: @@ -116,12 +120,14 @@ def main( biotype = "faa" if accession: - print("\nExtracting " + accession, flush=True) - hmm_path = runextractHmmAcc(hmm_path, accession, out_dir + accession + ".hmm") - if os.path.isfile(accession): + logger.debug("Extracting " + accession) + hmm_path = run_extract_hmm_acc( + hmm_path, accession, out_dir + accession + ".hmm" + ) + if Path(accession).is_file(): accession = [] else: - print("\nExtracting accessions", flush=True) + logger.info("Extracting accessions") with open(hmm_path, "r") as raw: hmm_data = raw.read() accession = [] @@ -131,18 +137,18 @@ def main( else: hmm_cpu = cpu hmmer_out = out_dir + "hmmer.out" - print("\nRunning " + binary, flush=True) - if runHmmer(fasta_path, hmm_path, hmmer_out, cpu=hmm_cpu, binary=binary): - eprint("\tERROR: " + binary + " failed", flush=True) + logger.debug("Running " + binary) + if run_hmmer(fasta_path, hmm_path, hmmer_out, cpu=hmm_cpu, binary=binary): + logger.error("" + binary + " failed") sys.exit(2) - print("\nParsing output", flush=True) + logger.info("Parsing output") hmm_data = run_extract_hmm( hmmer_out, top_hits, cov_threshold, evalue, not accession_search, accession ) output_res = parse_hmm_data(hmm_data) - print("\nCompiling fastas", flush=True) + logger.info("Compiling fastas") output_fas = run_acc2fa(db, biotype, output_res, subhit=subhit, cpu=cpu) return output_fas @@ -157,7 +163,7 @@ def cli(): parser.add_argument("--hmm", required=True, help="Input .hmm") parser.add_argument("-b", "--binary", required=True, help="{'hmmsearch', 'nhmmer'}") parser.add_argument( - "-d", "--mtdb", default=primaryDB(), help="MycoDB. DEFAULT: master" + "-d", "--mtdb", default=primary_db(), help="MycoDB. DEFAULT: master" ) parser.add_argument( "-q", "--query", help="Query [acc if -a] from .hmm, or new line delimited file" @@ -180,19 +186,20 @@ def cli(): parser.add_argument("-o", "--output", help="Output directory") parser.add_argument("--cpu", default=1, type=int) args = parser.parse_args() + setup_logging(verbose=getattr(args, "verbose", False)) if args.binary not in {"hmmsearch", "nhmmer"}: - eprint("\nERROR: invalid --binary", flush=True) + logger.error("invalid --binary") sys.exit(1) - findExecs([args.binary], exit={args.binary}) + find_execs([args.binary], exit={args.binary}) if args.output: out_dir = args.output else: date = datetime.datetime.today().strftime("%Y%m%d") - out_dir = os.getcwd() + "/" + date + "_fa2hmm2fa/" - if not os.path.isdir(out_dir): - os.mkdir(out_dir) + out_dir = str(Path.cwd()) + "/homology_fasta_" + date + "/" + if not Path(out_dir).is_dir(): + Path(out_dir).mkdir() out_dir = format_path(out_dir) if args.evalue: @@ -234,10 +241,8 @@ def cli(): accession_search=args.accession, subhit=not args.whole, ) - fastaname = re.sub( - r"\.fa[^\.]*$", "", os.path.basename(os.path.abspath(args.fasta)) - ) - outputFas(output_fas, out_dir, fastaname) + fastaname = re.sub(r"\.fa[^\.]*$", "", Path(os.path.abspath(args.fasta)).name) + output_fas(output_fas, out_dir, fastaname) outro(start_time) diff --git a/mycotools/jgiDwnld.py b/mycotools/jgiDwnld.py deleted file mode 100755 index 63fb703..0000000 --- a/mycotools/jgiDwnld.py +++ /dev/null @@ -1,831 +0,0 @@ -#! /usr/bin/env python3 -""" -PLEASE respect JGI's ping time limits. I've tuned it to respect their -unannounced limit. - - -NEED to remove gff v gff3 option -""" - -import os -import re -import sys -import time -import getpass -import argparse -import subprocess -import pandas as pd -import xml.etree.ElementTree as ET -from tqdm import tqdm -from mycotools.lib.kontools import eprint, format_path, outro, intro -from mycotools.lib.dbtools import loginCheck - - -def jgi_login(user, pwd): - """Login via JGI's prescribed method by creating a cookie cache and - downloading JGI's sign-in file.""" - - null = os.path.expanduser("~/.nulljgi_dwnld") - - login_cmd = subprocess.call( - [ - "curl", - "https://signon.jgi.doe.gov/signon/create", - "--data-urlencode", - "login=" + str(user), - "--data-urlencode", - "password=" + str(pwd), - "-c", - "cookies", - "-o", - null, - ], - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - ) - - return login_cmd - - -def dwnld_xml(output, ome, max_tempts=2): - attempts = 0 - while not os.path.isfile(f"{output}/{ome}.xml") and attempts < max_tempts: - attempts += 1 - xml_cmd = subprocess.call( - [ - "curl", - "https://genome.jgi.doe.gov/portal/ext-api/downloads/get-directory?organism=" - + str(ome), - "-b", - "cookies", - "-o", - f"{output}/{ome}.xml", - ], - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - ) - if xml_cmd != 0: - print(f"\tERROR: {ome} xml curl error: {xml_cmd}", flush=True) - if not os.path.isfile(f"{output}/{ome}.xml"): - return -1 - else: - return xml_cmd - - -def retrieve_xml(ome, output): - """Retrieve JGI xml file tree. First check if it already exists, if not then - download it using JGI's prescribed method. Then open the xml and check for - the common 'Portal does not exist' error. If so, report.""" - - if os.path.exists(output + "/" + str(ome) + ".xml"): - with open(output + "/" + str(ome) + ".xml", "r") as xml_raw: - xml_data = xml_raw.read() - if xml_data == "Portal does not exist": - print("\tERROR: `" + ome + " not in JGIs `organism` database", flush=True) - xml_cmd = 1 - os.remove(output + "/" + ome + ".xml") - elif not xml_data: - xml_cmd = None - os.remove(f"{output}/{ome}.xml") - else: - xml_cmd = -1 - else: - xml_cmd = dwnld_xml(output, ome) - - if xml_cmd == 0: - with open(output + "/" + ome + ".xml", "r") as xml_raw: - xml_data = xml_raw.read() - if xml_data == "Portal does not exist": - print("\tERROR: `" + ome + " not in JGIs `organism` database", flush=True) - xml_cmd = 1 - os.remove(output + "/" + ome + ".xml") - elif not xml_data: - xml_cmd = None - os.remove(f"{output}/{ome}.xml") - - return xml_cmd - - -def parse_xml(ft, xml_file, masked=False, forbidden={}, filtered=True): - """Parse the XML data to obtain the file types of interest based on - predefined hashes that contain the known subdirectories associated with JGI - organization""" - - if ft == "fna": - if masked: - ft += "$masked" - else: - ft += "$unmasked" - - # set the initial hashes for the XML hierarchy - relate file types to their - # hierarchy structure - ft2xt = { - "fna$masked": {"assembly"}, - "fna$unmasked": {"assembly"}, - "gff": {"annotation"}, - "gff3": {"annotation"}, - "transcripts": {"annotation"}, - "est": {"ests and est clusters", "transcriptome"}, - } - ft2fh = { - "fna$masked": ["genome assembly (masked)", "assembled scaffolds (masked)"], - "fna$unmasked": [ - "assembled scaffolds (unmasked)", - "genome assembly (unmasked)", - ], - "gff": ["genes"], - "gff3": ["genes"], - "transcripts": ["transcripts"], - "est": ["ests", "transcriptome assembly"], - } - - ft2fn = { - "fna$masked": ["masked", "Genome Assembly (masked)"], - "fna$unmasked": [ - "AssembledScaffolds", - "scaffolds", - "Genome Assembly (unmasked)", - "AssemblyScaffolds", - ], - "gff3": ["GeneCatalog", "FilteredModels"], - "transcripts": ["transcripts"], - "est": ["EST"], - } - ft2fe = { - "fna$masked": {"fasta", "fa", "fna", "fsa"}, - "fna$unmasked": {"fasta", "fa", "fna", "fsa"}, - "gff": {"gff", "gff3"}, - "gff3": {"gff", "gff3"}, - "transcripts": {"fa", "fasta", "fna", "fsa"}, - "est": {"fa", "fasta", "fna", "fsa"}, - } - - url, md5, filename = None, False, None - - # parse the XML file - tree = ET.parse(xml_file) - root = tree.getroot() - flip = True - org_name = None - has_flipped = False - attempt = 0 - - # flip is a way to rerun the loop if the file type changes (e.g. from - # masked to unmasked); parse through the XML hiearchy in accord with the - # hashes established above - while flip and attempt < 10: - attempt += 1 - for child in root: - # conserved subdirectory we need - if "Files" == child.attrib["name"]: - for chil1 in child: - # does this subdirectory match what we need for our - # filetype? - if chil1.attrib["name"].lower() in ft2xt[ft]: - # we only want the filtered models for annotations/RNA - if ft in {"gff3", "gff", "transcripts", "est"}: - for chil2 in chil1: - if ( - chil2.attrib["name"] - .lower() - .startswith("filtered models (") - ): - chil1 = chil2 - break - # continue parsing toward the files of interest - for chil2 in chil1: - if any( - x == chil2.attrib["name"].lower() for x in ft2fh[ft] - ): - for chil3 in chil2: - t_url = chil3.attrib["url"] - try: - org_name = chil3.attrib["label"] - except KeyError: - pass - # we want to avoid tape files as we cannot - # download them readily - if ( - "get_tape_file" not in t_url - and t_url not in forbidden - ): - if all( - x not in chil3.attrib["filename"] - for x in ft2fn[ft] - ): - continue - file_ext_srch = re.search( - r"\.([^\.]+)$", chil3.attrib["filename"] - ) - if file_ext_srch is not None: - file_ext = file_ext_srch[1] - if file_ext == "gz": - file_ext_srch = re.search( - r"\.([^\.]+)\.gz$", - chil3.attrib["filename"], - ) - if file_ext_srch is not None: - file_ext = file_ext_srch[1] - if file_ext not in ft2fe[ft]: - continue - else: - continue - - url = chil3.attrib["url"] - filename = chil3.attrib["filename"] - # all requirements satisfied - try: - md5 = chil3.attrib["md5"] - break - # continue on to find an md5, or omit - # if exhaustively searched - except KeyError: - pass - - # if the unmasked genome is not present, then query for the masked and - # vice versa - if not url and ft == "fna$masked": - ft = "fna$unmasked" - if not has_flipped: - flip = True - else: - flip = False - elif not url and ft == "fna$unmasked": - ft = "fna$masked" - if not has_flipped: - flip = True - else: - flip = False - else: - break - - return filename, url, md5, org_name - - -def handle_redirect_307( - dwnld_data, dwnld, dwnld_url, file_type, xml_file, masked, url, urls, spacer -): - """Handle a redirection error by identifying a new file URL to download - from, or return the original if none exist""" - print( - spacer + "\t" + dwnld + " link has moved. " + "Trying a different link.", - flush=True, - ) - filename, n_url, dwnld_md5, t_org_name = parse_xml( - file_type, xml_file, masked=masked, forbidden={url}.union(urls) - ) - - if n_url: - url = n_url - dwnld_url = prefix + url.replace("&", "&") - dwnld = f"{output}{file_type}/{os.path.basename(dwnld_url)}" - return url, dwnld_url, dwnld, {url}.union(urls), t_org_name - - -def no_md5_checks(dwnld, md5, spacer): - """If there is no MD5, simply check the file has content in it""" - check_size = subprocess.run(["wc", "-l", dwnld], stdout=subprocess.PIPE) - check_size_res = check_size.stdout.decode("utf-8") - print(spacer + "\t\tFile exists - no md5 to check.", flush=True) - check_size_find = re.search(r"\d+", check_size_res) - size = check_size_find[0] - if int(size) < 10: - print(spacer + "\tInvalid file size.", flush=True) - else: - md5 = None - return md5 - - -def jgi_dwnld(ome, file_type, output, masked=True, spacer="\t"): - """Download JGI files. For each type of file, use regular expressions to - gather the URL from the file. Grab the md5checksum if possible from the - xml as well. Create arbitrary values for md5 and curl_cmd. If the download file - already exists, then run an md5 checksum if an md5 value exists in the xml. - - If the md5 does not match then open and check for typical errors. If those - errors exist, obtain the alternative download URL from the xml. If the md5 - matches, pass through the rest of the function. - - If there is no download md5, then check to see if the file is greater than 10 - lines as a proxy to make sure that the file isn't empty/blatantly wrong. If it - passes this test, change the values to not enter the while loop at the end of - the function. - - The while loop following the file exists error allows for 3 attempts. If a file - is downloaded it will check its md5 using the xml reference. If it fails, it will - proceed via the error checking above, wait a minute, and reattempt. If there is no - download md5 from the xml it will proceed via the error checking above as well. If it - passes the line count check, it will exit the loop - otherwise, it will attempt to - gather a new URL and restart after another minute wait. This is an unfortunate - circumnavigation of JGI's cryptic maximum ping / time before booting.""" - - # prepare data structures - prefix = "https://genome.jgi.doe.gov" - xml_file = f"{output}xml/{ome}.xml" - preexisting, check = False, 1 - - # stop gap for legacy input - if file_type == "gff": - file_type += "3" - - ran_dwnld = False - org_name = None - - # acquire the filename, URL, and MD5 from the xml for the file type of - # interest - filename, url, dwnld_md5, t_org_name = parse_xml(file_type, xml_file, masked=masked) - if t_org_name: - org_name = t_org_name - if not dwnld_md5: - dwnld_md5 = None - - # if there is a URL present, begin the downloading process - if url: - f_urls = {url} - md5 = False - attempt = 0 - curl_cmd = 420 - - dwnld_url = prefix + url.replace("&", "&") - - dwnld = f"{output}{file_type}/{os.path.basename(dwnld_url)}" - unzip_dwnld = re.sub(r"\.gz$", "", dwnld) - # assume unzipped downloads have passed the checks - if os.path.isfile(unzip_dwnld): - md5 = dwnld_md5 - curl_cmd = 0 - check = unzip_dwnld - preexisting = True - - # if the file currently exists, then check its MD5 - elif os.path.exists(dwnld): - if dwnld_md5: - md5_cmd = subprocess.run( - ["md5sum", dwnld], stdout=subprocess.PIPE, stderr=subprocess.PIPE - ) - md5_res = md5_cmd.stdout.decode("utf-8") - md5_find = re.search(r"\w+", md5_res) - md5 = md5_find[0] - - if md5 == dwnld_md5: - curl_cmd = 0 - check = dwnld - preexisting = True - # if the MD5 does not equal the download MD5 then check the file - else: - while True: - try: - with open(dwnld, "r") as dwnld_data_raw: - dwnld_data = dwnld_data_raw.read() - if re.search("307 Temporary Redirect", dwnld_data): - url, dwnld_url, dwnld, f_urls, t_org_name = ( - handle_redirect_307( - dwnld_data, - dwnld, - dwnld_url, - file_type, - xml_file, - masked, - url, - f_urls, - spacer, - ) - ) - if t_org_name: - org_name = t_org_name - break - except FileNotFoundError: - md5 = False - break - except UnicodeDecodeError: - if not dwnld_md5: - md5 = no_md5_checks(dwnld, md5, spacer) - if not md5: - preexisting = True - check = dwnld - else: - print(spacer + "\tmd5 does not match.", flush=True) - break - - # while the MD5 doesn't match, or there is a curl error, try up to 3 - # times to download the file - while md5 != dwnld_md5 and curl_cmd != 0 and attempt < 3: - attempt += 1 - curl_cmd = subprocess.call( - ["curl", dwnld_url, "-b", "cookies", "-o", dwnld], - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - ) - ran_dwnld = True - - if curl_cmd == 0: - check = dwnld - - # acquire the MD5 - if dwnld_md5: - md5_cmd = subprocess.run( - ["md5sum", dwnld], - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - ) - md5_res = md5_cmd.stdout.decode("utf-8") - md5_find = re.search(r"\w+", md5_res) - try: - md5 = md5_find[0] - except TypeError: - md5 = False - if not os.path.isfile(dwnld): - attempt += 1 - continue - # if there is no MD5 attempt the crude file check - if not dwnld_md5: - try: - with open(dwnld, "r") as dwnld_data_raw: - dwnld_data = dwnld_data_raw.read() - if re.search("307 Temporary Redirect", dwnld_data): - t_url, dwnld_url, dwnld, f_urls, t_org_name = ( - handle_redirect_307( - dwnld_data, - dwnld, - dwnld_url, - file_type, - xml_file, - masked, - url, - f_urls, - spacer, - ) - ) - if t_org_name: - org_name = t_org_name - - if t_url == url: - print(spacer + "\t\tNo valid alternative", flush=True) - attempt = 4 - break - else: - url = t_url - else: - md5 = no_md5_checks(dwnld, md5, spacer) - if not md5: - break - except FileNotFoundError: - pass - except UnicodeDecodeError: - pass - # this is slow, and was arbitrarily set to not overping JGI - time.sleep(60) - - # the download may have failed, so prepare to retry - elif md5 != dwnld_md5 and attempt == 1: - print( - f"{spacer}\tERROR: md5 does not match JGI. " - + f"Attempt {attempt}", - flush=True, - ) - curl_cmd = -1 - check = 2 - while True: - try: - with open(dwnld, "r") as dwnld_data_raw: - dwnld_data = dwnld_data_raw.read() - if re.search("307 Temporary Redirect", dwnld_data): - t_url, dwnld_url, dwnld, f_urls, t_org_name = ( - handle_redirect_307( - dwnld_data, - dwnld, - dwnld_url, - file_type, - xml_file, - masked, - url, - f_urls, - spacer, - ) - ) - if t_org_name: - org_name = t_org_name - if t_url == url: - print( - spacer + "\t\tNo valid alternative", flush=True - ) - attempt = 4 - break - break - except FileNotFoundError: - pass - break - except UnicodeDecodeError: - pass - break - time.sleep(60) - # if there are two fails, attempt a new URL - elif md5 != dwnld_md5 and attempt == 2: - print( - f"{spacer}\tERROR: md5 does not match JGI. Attempt {attempt}", - flush=True, - ) - curl_cmd = -1 - filename, n_url, dwnld_md5, t_org_name = parse_xml( - file_type, xml_file, masked=masked, forbidden=f_urls - ) - if t_org_name: - org_name = t_org_name - - if n_url: - url = n_url - dwnld_url = prefix + url.replace("&", "&") - f_ulrs = {url}.union(f_urls) - dwnld = f"{output}{file_type}/{os.path.basename(dwnld_url)}" - time.sleep(60) - check = 2 - elif md5 != dwnld_md5: - print( - f"{spacer}\tERROR: md5 does not match JGI. Attempt {attempt}", - flush=True, - ) - check = 2 - else: - print( - f"{spacer}\tERROR: Failed to retrieve {file_type}. `curl` error: " - + f"{curl_cmd}\n{spacer}\tAttempt {attempt}", - flush=True, - ) - check = 2 - - # three strikes and the file is out - if attempt == 3: - if md5 != dwnld_md5: - print( - spacer + "\tExcluding from database - potential failure", flush=True - ) - curl_cmd = 0 - if curl_cmd != 0: - print(spacer + "\tFile failed to download", flush=True) - - return check, preexisting, file_type, ran_dwnld, org_name - - -def main( - df, - output, - user, - pwd, - assembly=True, - proteome=False, - gff3=True, - transcript=False, - est=False, - masked=True, - spacer="\t", -): - # pd.options.mode.chained_assignment = None # default='warn' - if not "assembly_acc" in df.columns: - if len(df.columns) != 1: - eprint( - "\nInvalid input. No assembly_acc column and more than one column.", - flush=True, - ) - else: - ome_col = list(df.columns)[0] - else: - ome_col = "assembly_acc" - - eprint(spacer + "Logging into JGI", flush=True) - login_attempt = 0 - while jgi_login(user, pwd) != 0 and login_attempt < 5: - eprint( - spacer + "\tJGI Login Failed. Attempt: " + str(login_attempt), flush=True - ) - time.sleep(5) - login_attempt += 1 - if login_attempt == 3: - eprint(spacer + "\tERROR: Failed 3 login attempts.", flush=True) - sys.exit(100) - - if not os.path.exists(output + "/xml"): - os.mkdir(output + "/xml") - # perhaps add a counter here, but one that checks if it is actually querying jgi - print("\nRetrieving `xml` directories", flush=True) - ome_set, count = set(), 0 - for i, row in tqdm(df.iterrows(), total=len(df)): - error_check, attempt = True, 0 - while error_check != -1 and attempt < 3: - attempt += 1 - error_check = retrieve_xml(row[ome_col], output + "/xml") - if error_check is None: - time.sleep(1) - continue - # elif error_check > 0: - # ome_set.add(row[ome_col]) - elif error_check != -1: - time.sleep(0.3) - if error_check != -1: - eprint(f"{spacer}\t{row[ome_col]} failed to retrieve XML", flush=True) - ome_set.add(row[ome_col]) - - eprint( - f"{spacer}Downloading {len(df)} JGI files\n\t" + "Maximum rate: 1 file/min", - flush=True, - ) - - dwnlds = [] - if assembly: - dwnlds.append("fna") - if proteome: - dwnlds.append("faa") - if gff3: - dwnlds.append("gff3") - if transcript: - dwnlds.append("transcript") - # dwnlds.append( 'AllTranscript' ) - if est: - dwnlds.append("est") - - for typ in dwnlds: - if not os.path.isdir(output + "/" + typ): - os.mkdir(output + "/" + typ) - - preexisting, ran_dwnld = True, False - for i, row in df.iterrows(): - ome = row[ome_col] - if ran_dwnld: - time.sleep(60) - if ome not in ome_set: - jgi_login(user, pwd) - if "ome" in row.keys(): - eprint(spacer + row["ome"] + "\t" + ome, flush=True) - else: - eprint(spacer + ome, flush=True) - for typ in dwnlds: - check, preexisting, new_typ, ran_dwnld, org_name = jgi_dwnld( - ome, typ, output, masked=masked, spacer=spacer - ) - if type(check) != int: - df.at[i, new_typ + "_path"] = ( - output + "/" + new_typ + "/" + os.path.basename(check) - ) - check = os.path.basename(os.path.abspath(check)) - if org_name: - org_d = org_name.split() - genus = org_d[0] - if len(org_d) > 1: - sp = org_d[1] - else: - sp = "sp." - if len(org_d) > 2: - strain = "".join(org_d[2:]) - else: - strain = "" - else: - genus, sp, strain = "", "", "" - df.at[i, "genus"] = genus - df.at[i, "species"] = sp - df.at[i, "strain"] = strain - elif type(check) == int: - ome_set.add(row[ome_col]) - eprint( - spacer + "\t" + new_typ + ": exit status " + str(check), flush=True - ) - else: - eprint(spacer + ome + " failed.", flush=True) - - if os.path.exists("cookies"): - os.remove("cookies") - if os.path.exists(os.path.expanduser("~/.null")): - os.remove(os.path.expanduser("~/.null")) - - if "gff3" in df.columns: - del df["gff3"] - if "faa" in df.columns: - del df["faa"] - if "fna" in df.columns: - del df["fna"] - - return df, ome_set - - -def cli(): - - parser = argparse.ArgumentParser( - description="Imports table/database with JGI `assembly_acc` column and downloads assembly, proteome, gff, " - + "and/or gff3. This script supports rerunning/continuing previous runs in the same directory. " - + "JGI has stringent, nondefined ping limits, so file downloads are limited to 1 per minute." - ) - parser.add_argument( - "-i", - "--input", - required=True, - help="Genome code or table with `assembly_acc` column of JGI ome codes", - ) - parser.add_argument( - "-a", - "--assembly", - default=False, - action="store_true", - help="Download assembly fastas", - ) - parser.add_argument( - "-p", - "--proteome", - default=False, - action="store_true", - help="Download proteome fastas", - ) - # parser.add_argument( '-g', '--gff', default = False, action = 'store_true', \ - # help = 'Download gffs.' ) - parser.add_argument( - "-g", "--gff", default=False, action="store_true", help="Download gff3s" - ) - parser.add_argument( - "-t", - "--transcript", - default=False, - action="store_true", - help="Download transcripts fastas", - ) - parser.add_argument( - "-e", "--est", default=False, action="store_true", help="Download EST fastas" - ) - parser.add_argument( - "--nonmasked", - default=False, - action="store_true", - help="[-a] Download nonmasked assemblies", - ) - parser.add_argument("-o", "--output", default=os.getcwd(), help="Output dir") - args = parser.parse_args() - - if args.nonmasked: - args.assembly = True - - if ( - not args.assembly - and not args.proteome - and not args.transcript - and not args.est - and not args.gff - ): - eprint("\nERROR: You must choose at least one download option.", flush=True) - - ncbi_email, ncbi_api, user, pwd = loginCheck(ncbi=False) - # user = input( 'JGI username: ' ) - # pwd = getpass.getpass( prompt='JGI Login Password: ' ) - - args_dict = { - "JGI Table": args.input, - "Assemblies": args.assembly, - "RepeatMasked": not args.nonmasked, - "Proteomes": args.proteome, - ".gff3's": args.gff, - "Transcripts": args.transcript, - "EST": args.est, - } - - start_time = intro("Download JGI files", args_dict) - eprint( - "\nWARNING: This script does NOT account for use-restricted data. " - + "It is user responsibility to determine use restriction status " - + "in accord with the MycoCosm terms and conditions: " - + "https://jgi.doe.gov/user-programs/pmo-overview/policies/legacy-data-policies/", - flush=True, - ) - eprint(flush=True) - - if os.path.isfile(args.input): - with open(args.input, "r") as raw: - for line in raw: - if "assembly_acc" in line.rstrip().split("\t"): - df = pd.read_csv(args.input, sep="\t", index_col=None) - else: - df = pd.read_csv(args.input, sep="\t", header=None) - break - else: - in_data = args.input.replace('"', "").replace("'", "").replace(",", " ").split() - df = pd.DataFrame({"assembly_acc": in_data}) - - output = format_path(args.output) - - jgi_df, ome_set = main( - df, - output, - user, - pwd, - args.assembly, - args.proteome, - args.gff, - args.transcript, - args.est, - not args.nonmasked, - spacer="", - ) - jgi_df = jgi_df.rename(columns={"assembly_acc": "#assembly_acc"}) - jgi_df["source"] = "jgi" - jgi_df["restriction"] = "no" - jgi_df.to_csv(os.path.normpath(args.input) + ".predb.tsv", sep="\t", index=False) - - outro(start_time) - - -if __name__ == "__main__": - cli() diff --git a/mycotools/lib/biotools.py b/mycotools/lib/biotools.py index fdfbccf..4c8905f 100755 --- a/mycotools/lib/biotools.py +++ b/mycotools/lib/biotools.py @@ -3,10 +3,8 @@ # NEED to convert gff list to appropriate types import re -import sys from collections import defaultdict from itertools import chain -from mycotools.lib.kontools import eprint aa_weights = { "A": 89.1, @@ -323,49 +321,6 @@ def fa2dict_accs(fasta_input, accs=set()): # file_ = True): fasta_dict[seq_name]["sequence"] += data return fasta_dict - # if file_: - - -# with open(fasta_input, 'r') as fasta: -# str_fasta = '' -# for line in fasta: -# data = line.rstrip() -# if data.startswith('>'): -# str_fasta += '\n' + data.rstrip() + '\n' -# elif data: -# str_fasta += data -# concatenates the prep string with the prepped line from the fasta file -# str_fasta = str_fasta.lstrip() -# else: -# str_fasta = fasta_input - -# extracts 0) seq ID and description and 1) sequence -# extracted = re.findall( r'(^>[^\n]*)\n([^>]*)', str_fasta, re.M) - -# adds a new dictionary for each gene -# for val in extracted: -# step1 = val[0] -# step2 = re.search(r'^>([^ ]*)', step1) -# gene = step2[1] -# step3 = re.search(r' (.*)', step1) -# if step3: -# descrip = step3[1] -# else: -# descrip = '' - -# seq = val[1] -# if seq: -# if seq[-1] == '\n' or seq[-1] == '\r': -# seq = seq.rstrip() -# else: -# continue - -# prepares dictionaries for each gene with description, seq, rvcmpl_seq, and codons -# fasta_dict[gene] = {} -# if descrip != '\n': -# fasta_dict[gene]['description'] = descrip -# fasta_dict[gene]['sequence'] = seq - # truncates sequences based on inputted lenght def dnatrunc(fasta_dict, trunc_length): @@ -452,10 +407,353 @@ def calc_gc(gene): return GC_con +class GFFList(): + def __init__(self): + pass + + +# GFF strand column -> BioPython-style strand integer ('.'/'?' -> None) +_STRAND = {"+": 1, "-": -1} +# strand integer -> GFF strand column (anything else, e.g. None, serializes to '.') +_STRAND_REVERSE = {1: "+", -1: "-"} +# the placeholder characters GFF3 uses for an undefined numeric/strand column +_UNDEFINED = {".", "?", "", None} + +# GFF3 column-9 reserved characters and their percent-encodings; '%' is listed +# first so the escapes introduced below are not themselves re-encoded +_GFF_ATTR_ESCAPES = ( + ("%", "%25"), (";", "%3B"), ("=", "%3D"), ("&", "%26"), (",", "%2C"), + ("\t", "%09"), ("\n", "%0A"), ("\r", "%0D"), +) + +def _gff_escape(value: str) -> str: + """Percent-encode the GFF3 attribute-reserved characters in ``value``""" + for char, code in _GFF_ATTR_ESCAPES: + value = value.replace(char, code) + return value + + +def _gff_unescape(value: str) -> str: + """Percent-decode the GFF3 attribute-reserved characters in ``value``. + + Iterates in reverse so '%25' -> '%' is applied last, preventing an already + decoded literal from being decoded a second time.""" + for char, code in reversed(_GFF_ATTR_ESCAPES): + value = value.replace(code, char) + return value + + +def _format_gff_attributes(attributes: dict, field_delimiter: str = ";", value_delimiter: str = "=") -> str: + """Serialize a ``{key: value}`` attribute dict to a GFF3 column-9 string. + + Reserved characters in keys and values are percent-encoded; an empty dict + yields ``.`` (the GFF3 "no attributes" placeholder).""" + if not attributes: + return "." + return field_delimiter.join( + f"{_gff_escape(str(key))}{value_delimiter}{_gff_escape(str(value))}" + for key, value in attributes.items() + ) + + +def _parse_gff_attributes(attributes: str, field_delimiter: str = ";", value_delimiter: str = "=") -> dict: + """Parse a GFF3 column-9 string into a ``{key: value}`` dict, inverting + ``_format_gff_attributes``. + + Fields are split on ``field_delimiter`` and each into key/value on the first + ``value_delimiter``; reserved characters are percent-decoded. A valueless + field maps to ``""``; the '.' placeholder (or an empty string) yields an + empty dict.""" + parsed = {} + if not attributes or attributes in _UNDEFINED: + return parsed + for field in attributes.split(field_delimiter): + field = field.strip() + if not field: + continue + key, sep, value = field.partition(value_delimiter) + parsed[_gff_unescape(key.strip())] = _gff_unescape(value.strip()) if sep else "" + return parsed + + +def _as_int(value, name: str): + """Coerce a GFF numeric column to ``int``, mapping the undefined placeholder + ('.') to ``None``""" + if value in _UNDEFINED: + return None + try: + return int(value) + except (TypeError, ValueError): + raise ValueError(f"{name} must be an integer, got {value!r}") + + +def _as_strand(value): + """Coerce a GFF strand column ('+'/'-') or a signed integer to 1/-1, with the + undefined placeholder ('.'/'?') mapping to ``None``""" + if value in _UNDEFINED: + return None + if value in _STRAND: + return _STRAND[value] + if value in (1, -1): + return value + raise ValueError(f"strand must be one of '+'/'-'/1/-1, got {value!r}") + + +class Feature: + """A single cross-annotation feature + `start` and `end` are sorted by smallest to largest and are expected to be 0-based half-open upon input""" + + def __init__(self, fid=None, pid=None, seqid=None, source=None, type=None, start=None, end=None, score=None, strand=None, phase=None, + attributes=None, descendants: list = None, parent: "Feature" = None, sequence: str = "", origin_sequence: str = "", + ingest: bool = False + ): + # feature ID + self.fid = fid + # parent ID + self.pid = pid + self.seqid = seqid + self.source = source + self.type = type + put_start = _as_int(start, "start") + put_end = _as_int(end, "end") + # enforce start < end + self.start, self.end = sorted((put_start, put_end)) + self.score = _as_int(score, "score") + self.strand = _as_strand(strand) + self.phase = _as_int(phase, "phase") + self.parent = parent + + # sentinel defaults: a shared mutable default would let every Feature + # alias one dict/list, collapsing the hierarchy built by group_features + self.attributes = {} if attributes is None else attributes + self.descendants = [] if descendants is None else descendants + + # derive sequence from the full contiguous sequence provided + if origin_sequence: + # 0-BASED, HALF-OPEN + sequence = origin_sequence[self.start:self.end] + + # demand that the provided sequence abide by the coordinates + if sequence: + if len(sequence) < self.end - self.start: + raise ValueError(f"sequence length deviates from coordinate length") + else: + self.sequence = sequence + else: + self.sequence = "" + + if ingest: + self._ingest() + + if not self.fid: + raise AttributeError('No ID obtained for feature') + + + def _ingest(self): + """Normalize ``attributes`` and derive ``fid``/``pid`` from them. + + A string ``attributes`` (a raw GFF3 column-9 field) is parsed into a + dict; ``fid`` is then set from the first present of ID/Id/id and ``pid`` + from the first of Parent/parent. An attribute hit overrides the + constructor value; a miss leaves it untouched.""" + if isinstance(self.attributes, str): + self.attributes = _parse_gff_attributes(self.attributes) + if not self.fid: + self.fid = _attr_get(self.attributes, ("ID", "Id", "id")) + if not self.pid: + self.pid = _attr_get(self.attributes, ("Parent", "parent")) + + + def _to_gff_line(self) -> str: + """Serialize this feature (without its descendants) to one GFF3 line. + + Coordinates are converted from the internal 0-based, half-open + representation back to GFF3's 1-based, both-inclusive columns; an + undefined ``start``/``end``/``score``/``strand``/``phase`` renders as the + '.' placeholder.""" + start = "." if self.start is None else str(self.start + 1) + end = "." if self.end is None else str(self.end) + score = "." if self.score is None else str(self.score) + strand = _STRAND_REVERSE.get(self.strand, ".") + phase = "." if self.phase is None else str(self.phase) + id_dict = {"ID": self.fid} + if self.pid: + id_dict["Parent"] = self.pid + return "\t".join(( + self.seqid, self.source, self.type, start, end, + score, strand, phase, _format_gff_attributes({**id_dict, **self.attributes}), + )) + + def to_gff(self) -> str: + """Serialize this feature and all of its descendants to a GFF3 string. + + Lines are emitted depth-first, each parent before its children, joined + by newlines (no trailing newline).""" + lines = [self._to_gff_line()] + for descendant in self.descendants: + lines.append(descendant.to_gff()) + return "\n".join(lines) + + +def _attr_get(attributes: dict, keys): + """Return the value of the first present attribute key (accepts case + variants, e.g. ``ID``/``id`` or ``Parent``/``parent``), or ``None``""" + for key in keys: + if key in attributes: + return attributes[key] + return None + + +def group_features(features: list, parent_ids: list = ["Parent", "parent"], ids: list = ["ID", "Id", "id"]): + """Hierarchically group features based on Parent <-> ID relationships""" + id2feature = {} + for feature in features: + fid = _attr_get(feature.attributes, ids) + if fid in id2feature: + raise KeyError(f"{fid} is depicted in multiple Features") + id2feature[fid] = feature + + feature_dict = defaultdict(list) + for fid, feature in id2feature.items(): + par_id = _attr_get(feature.attributes, parent_ids) + if par_id: + # link features together + id2feature[par_id].descendants.append(feature) + feature.parent = id2feature[par_id] + feature_dict[feature.seqid].append(feature) + + return feature_dict + + +# GFF `type` tokens grouped into the canonical feature classes FeatureCol +# exposes. RNA is a category: any type ending in "rna"/"transcript" (mRNA, +# tRNA, ncRNA, primary_transcript, ...) resolves to the "rna" class. +_FEATURE_CLASSES = ("gene", "rna", "cds", "exon") + + +class FeatureCol: + """An ordered collection of Features linked by Parent <-> ID relationships. + + On construction the supplied Features are hierarchically grouped (see + `group_features`), populating each Feature's `parent`/`descendants`. + Features of a canonical class are then reachable both as attributes + (`.genes`, `.rnas`, `.cds`, `.exons`) and by key (`fl["gene"]`, + `fl["mRNA"]`); both resolve to the same lists, extracted once by + `_extract_types`.""" + +# NEED ID keying +# NEED general buckets for features + def __init__(self, features: list = None, group: bool = True): + self.features = list(features) if features is not None else [] + if group: + # wire up parent/descendant links in place + group_features(self.features) + self._extract_types() + + @staticmethod + def _class_of(ftype): + """Return the canonical class ('gene'/'rna'/'cds'/'exon') for a GFF + `type`, or None if it is not one of them.""" + if not ftype: + return None + t = ftype.lower() + if t.endswith("rna") or t.endswith("transcript"): + return "rna" + return None + + def _extract_types(self): + """Bucket features by canonical class and store each list as an + attribute (`self.genes`, `self.rnas`, `self.cds`, `self.exons`).""" + buckets = {name: [] for name in _FEATURE_CLASSES} + for feature in self.features: + cls = self._class_of(feature.type) + if cls is not None: + buckets[cls].append(feature) + self._buckets = buckets + self.genes = buckets["gene"] + self.rnas = buckets["rna"] + self.cds = buckets["cds"] + self.exons = buckets["exon"] + + def sort(self): + """Order features by contig (seqid), then parent-before-descendant, + then start coordinate; returns self. + + Root features (no parent within this collection) are grouped by seqid + and ordered by start, then each is emitted immediately before its + descendants. Every descendant list is sorted by start in place, so the + depth-first walk here and `Feature.to_gff` share one sibling order. An + undefined ('.') seqid or start sorts to the end of its group.""" + def _start_key(feature): + return (feature.start is None, feature.start) + + for feature in self.features: + feature.descendants.sort(key=_start_key) + + def _walk(feature): + yield feature + for descendant in feature.descendants: + yield from _walk(descendant) + + roots = [f for f in self.features if f.parent is None] + roots.sort(key=lambda f: (f.seqid is None, f.seqid, _start_key(f))) + + ordered = [] + for root in roots: + ordered.extend(_walk(root)) + self.features = ordered + self._extract_types() + return self + + def roots(self): + """Return a new FeatureCol of only the root features (those with no + parent within this collection). + + Each root retains its existing `descendants`/`parent` wiring, so the + full hierarchy stays reachable through `.descendants`; grouping is + skipped to preserve those links rather than rebuild them.""" + return FeatureCol([f for f in self.features if f.parent is None], group=False) + + def _resolve_key(self, key: str): + """Map an access key to a canonical class: a class name ('gene'), its + plural ('genes'), or a raw GFF type ('mRNA').""" + k = key.lower() + if k in self._buckets: + return k + if k.endswith("s") and k[:-1] in self._buckets: + return k[:-1] + cls = self._class_of(k) + if cls is None: + raise KeyError(f"{key!r} is not a gene/RNA/CDS/exon feature class") + return cls + + def __getitem__(self, key): + """Key by feature class ('gene'/'rna'/'cds'/'exon', a plural, or a raw + GFF type such as 'mRNA') to get that class's list; index or slice into + the full feature list with an int or slice.""" + if isinstance(key, (int, slice)): + return self.features[key] + if isinstance(key, str): + return self._buckets[self._resolve_key(key)] + raise TypeError( + f"FeatureCol keys must be str, int, or slice, not {type(key).__name__}" + ) + + def __iter__(self): + return iter(self.features) + + def __len__(self): + return len(self.features) + + def __repr__(self): + return (f"FeatureCol({len(self.genes)} genes, {len(self.rnas)} RNAs, " + f"{len(self.cds)} CDS, {len(self.exons)} exons)") + + # need to change into a class def gff2list(gff_info, path=True, error=True): - gff_list_dict = [] + feature_list = [] data = [] if path: with open(gff_info, "r") as raw_gff: @@ -465,8 +763,6 @@ def gff2list(gff_info, path=True, error=True): data.append(d.split("\t")) elif d.startswith("##FASTA"): break - # data = [x.split('\t') for x in raw_gff.read().split('\n') \ - # if x and not x.startswith('#')] else: for line in gff_info.split("\n"): if not line.startswith("#") and line: @@ -474,23 +770,20 @@ def gff2list(gff_info, path=True, error=True): data.append(d) elif line.startswith("##FASTA"): break - - # data = [x.split('\t') for x in gff_info.split('\n') \ - # if x and not x.startswith('#')] try: for col_list in data: - gff_list_dict.append( - { - "seqid": col_list[0], - "source": col_list[1], - "type": col_list[2], - "start": int(col_list[3]), - "end": int(col_list[4]), - "score": col_list[5], - "strand": col_list[6], - "phase": col_list[7], - "attributes": col_list[8], - } + feature_list.append( + Feature ( + seqid = col_list[0], + source = col_list[1], + type = col_list[2], + start = int(col_list[3]), + end = int(col_list[4]), + score = col_list[5], + strand = col_list[6], + phase = col_list[7], + attributes = col_list[8], + ) ) except IndexError: raise IndexError( @@ -504,23 +797,26 @@ def gff2list(gff_info, path=True, error=True): str(col_list[4:6]) + " invalid integer " + "conversion: " + str(col_list) ) - return gff_list_dict + return FeatureCol(feature_list) def list2gff(gff_list, ver=3): + """Serialize to a GFF string, prefixed with a ``##gff-version`` header when + ``ver`` is truthy. + Accepts either a ``FeatureCol`` (each Feature emits one line in list order) + or the legacy list of column dicts from ``gff2list``.""" if ver: gff_str = "##gff-version " + str(ver) + "\n" else: gff_str = "" - for line in gff_list: - add_str = "\t".join(str(val) for val in list(line.values())) - gff_str += add_str + "\n" + for feature in gff_list: + gff_str += feature._to_gff_line() + "\n" return gff_str.rstrip() -def gff3Comps(source=None): +def gff3_comps(source=None): comps = {} comps["par"] = "(?:^|(?<=;))" + r'Parent=["\']?([^;\'"]+)' @@ -545,7 +841,7 @@ def gff3Comps(source=None): return comps -def gff2Comps(): +def gff2_comps(): comps = {} comps["id"] = r'name "([^"]+)"' @@ -559,7 +855,7 @@ def gff2Comps(): return comps -def gtfComps(): +def gtf_comps(): comps = {} comps["id"] = r'gene_id "?([^"]+)"?' @@ -570,7 +866,7 @@ def gtfComps(): return comps -def compileExon(gff): +def compile_exon(gff): exon_dict = {} diff --git a/mycotools/lib/dbtools.py b/mycotools/lib/dbtools.py index dc8793e..422f55c 100755 --- a/mycotools/lib/dbtools.py +++ b/mycotools/lib/dbtools.py @@ -2,11 +2,14 @@ # NEED to make login check more intuitive and easier for NCBI only +from __future__ import annotations + import os import re import sys import copy import json +import logging import time import base64 import urllib @@ -14,17 +17,22 @@ import zipfile import datetime import subprocess +import random from tqdm import tqdm from Bio import Entrez -from io import StringIO +from typing import Any, Dict, Iterable, Mapping, Optional, Union from collections import defaultdict from mycotools.lib.kontools import ( + atomic_write, collect_files, - eprint, format_path, read_json, write_json, ) +from mycotools.lib import mtdb_sql +from pathlib import Path + +logger = logging.getLogger(__name__) class mtdb(dict): @@ -50,24 +58,28 @@ class mtdb(dict): "gff3", ] - # NEED a detect index feature for adding dicts in with alternative indices - def __init__(self, db=None, index=None, add_paths=True): - self.columns = [ - "ome", + #: columns that may serve as the MTDB index (unique, hashable keys) + _index_columns = frozenset({"assembly_acc", "ome"}) + #: taxonomic ranks excluded when assimilating taxonomy dictionaries + _forbidden_tax_ranks = frozenset( + { + "no rank", + "subkingdom", "genus", "species", - "strain", - "taxonomy", - "version", - "source", - "biosample", - "assembly_acc", - "acquisition_date", - "published", - "fna", - "faa", - "gff3", - ] + "species group", + "varietas", + "forma", + } + ) + + # NEED a detect index feature for adding dicts in with alternative indices + def __init__( + self, + db: Union[str, Mapping[str, Any], "mtdb", None] = None, + index: Optional[str] = None, + add_paths: bool = True, + ): if not db: if not index: super().__init__({x: [] for x in self.columns}) @@ -90,17 +102,18 @@ def __init__(self, db=None, index=None, add_paths=True): ): super().__init__(db) else: - super().__init__(mtdb.db2df(self, db, add_paths=add_paths)) + super().__init__(self.db2df(db, add_paths=add_paths)) self.index = index - def mtdb2pd(self): + def mtdb2pd(self) -> "pd.DataFrame": import pandas as pd copy_mtdb = copy.deepcopy(self) copy_mtdb = copy_mtdb.reset_index() return pd.DataFrame(copy_mtdb) # assume pd is imported - def pd2mtdb(df): # legacy integration + @staticmethod + def pd2mtdb(df: "pd.DataFrame") -> "mtdb": # legacy integration df = df.fillna("") for i, row in df.iterrows(): if not row["gff3"]: @@ -112,132 +125,171 @@ def pd2mtdb(df): # legacy integration db = mtdb({x: list(df[x]) for x in mtdb.columns}) return db - def db2df(self, db_path, add_paths=True): - df = defaultdict(list) - if os.stat(db_path).st_size == 0: + def db2df(self, db_path: str, add_paths: bool = True) -> Dict[str, list]: + """Read a database from disk into the column dict this class holds. + + Dispatches on the file itself, so a SQLite primary database and a + tab-delimited `.mtdb` interchange file are interchangeable everywhere a + path is accepted.""" + db_path = format_path(db_path) + if mtdb_sql.is_sqlite(db_path): + return mtdb_sql.read_db(db_path, add_paths=add_paths) + return self._read_flat(db_path, add_paths=add_paths) + + @classmethod + def from_string(cls, data: str, add_paths: bool = True) -> "mtdb": + """Build an MTDB from the text of a `.mtdb` file. + + This is the stdin path -- `mtdb extract -d -` and anything else piping a + database between tools.""" + db = cls() + lines = [ + x.rstrip().split("\t") + for x in data.splitlines() + if not x.startswith("#") and x.rstrip() + ] + parsed = db._parse_rows(lines, "", add_paths=add_paths) + db.clear() + db.update(parsed) + db.index = None + return db + + def _read_flat(self, db_path: str, add_paths: bool = True) -> Dict[str, list]: + """Read a tab-delimited `.mtdb` interchange file.""" + if Path(db_path).stat().st_size == 0: return {x: [] for x in mtdb.columns} - with open(format_path(db_path), "r") as raw: + with open(db_path, "r") as raw: data = [ x.rstrip().split("\t") for x in raw if not x.startswith("#") and x.rstrip() ] + return self._parse_rows(data, db_path, add_paths=add_paths) + + def _parse_rows( + self, data: "list[list[str]]", origin: str, add_paths: bool = True + ) -> Dict[str, list]: + """Turn split `.mtdb` fields into the column dict, validating arity.""" + df = defaultdict(list) columns = self.columns - for entry in data: + n_columns = len(columns) + for line_no, entry in enumerate(data, 1): + if len(entry) > n_columns: + raise ValueError( + f"{origin} line {line_no}: {len(entry)} fields, expected at " + f"most {n_columns}. Columns are {', '.join(columns)}" + ) [df[c].append("") for c in columns] # add a blank entry to each # column for i, d in enumerate(entry): df[columns[i]][-1] = d try: - df["taxonomy"][-1] = read_tax(df["taxonomy"][-1]) + df["taxonomy"][-1] = self.read_tax(df["taxonomy"][-1]) except json.decoder.JSONDecodeError: - print(df["taxonomy"][-1]) - sys.exit() + raise ValueError( + f"{origin} line {line_no}: malformed taxonomy " + f"{df['taxonomy'][-1]!r}" + ) df["taxonomy"][-1]["genus"] = df["genus"][-1] df["taxonomy"][-1]["species"] = df["genus"][-1] + " " + df["species"][-1] df["taxonomy"][-1]["strain"] = df["strain"][-1] if not add_paths: return df - try: - for i, ome in enumerate(df["ome"]): - if not df["fna"][i]: - if {"MYCOFNA", "MYCOFAA", "MYCOGFF3"}.difference( - set(os.environ.keys()) - ): - raise FileNotFoundError( - "You are not connected to a primary MTDB. " - + "Standalone databases need absolute paths" - ) - df["fna"][i] = os.environ["MYCOFNA"] + ome + ".fna" - df["faa"][i] = os.environ["MYCOFAA"] + ome + ".faa" - df["gff3"][i] = os.environ["MYCOGFF3"] + ome + ".gff3" - elif df["fna"][i] == ome + ".fna": - if {"MYCOFNA", "MYCOFAA", "MYCOGFF3"}.difference( - set(os.environ.keys()) - ): - raise FileNotFoundError( - "You are not connected to a primary MTDB. " - + "Standalone databases need absolute paths" - ) - df["fna"][i] = os.environ["MYCOFNA"] + ome + ".fna" - df["faa"][i] = os.environ["MYCOFAA"] + ome + ".faa" - df["gff3"][i] = os.environ["MYCOGFF3"] + ome + ".gff3" - except KeyError: - eprint( - "ERROR: MycotoolsDB not in path, cannot delineate biofile paths", - flush=True, + # read the data directories once: os.environ decodes the entire + # environment on each access, which otherwise dominates load time and + # makes it scale with the size of the user's shell environment + env = mtdb_sql.read_path_env() + needs_env = [ + i + for i, ome in enumerate(df["ome"]) + if not df["fna"][i] or df["fna"][i] == ome + ".fna" + ] + if not needs_env: + return df + if env is None: + raise FileNotFoundError( + "You are not connected to a primary MTDB. " + + "Standalone databases need absolute paths" ) + fna_dir, faa_dir, gff3_dir = env["MYCOFNA"], env["MYCOFAA"], env["MYCOGFF3"] + for i in needs_env: + ome = df["ome"][i] + df["fna"][i] = fna_dir + ome + ".fna" + df["faa"][i] = faa_dir + ome + ".faa" + df["gff3"][i] = gff3_dir + ome + ".gff3" return df - def df2db(self, db_path=None, headers=False, paths=False): - df = copy.copy(self) - df = df.reset_index() - output = mtdb( - { - k: v - for k, v in sorted(self.set_index("ome").items(), key=lambda x: x[0]) - }, - index="ome", - ) - # does this work if its not an inplace change - abb_paths = { - "faa": [os.environ["MYCOFAA"], ".faa"], - "fna": [os.environ["MYCOFNA"], ".fna"], - "gff3": [os.environ["MYCOGFF3"], ".gff3"], - } + def _export_rows(self, paths: bool = False) -> "list[list[str]]": + """Render this MTDB as the ordered field lists a `.mtdb` file holds. + + Row dicts are copied before the genome-level ranks are stripped from + their taxonomy: `set_index` shares the taxonomy dict with the caller, so + editing it in place would delete genus/species/strain out from under an + MTDB that is still in use.""" + env = mtdb_sql.read_path_env() + rows = [] + for ome, row in sorted(self.set_index("ome").items(), key=lambda x: x[0]): + row = dict(row) + if not paths and env is not None: + for file_type, var in ( + ("fna", "MYCOFNA"), + ("faa", "MYCOFAA"), + ("gff3", "MYCOGFF3"), + ): + default = env[var] + ome + "." + file_type + if str(row.get(file_type) or "") == default: + row[file_type] = "" # abbreviate when possible + taxonomy = row.get("taxonomy") + if isinstance(taxonomy, dict): + taxonomy = { + k: v + for k, v in taxonomy.items() + if k not in {"species", "genus", "strain"} + } + row["taxonomy"] = json.dumps(taxonomy) if taxonomy else "{}" + if not row.get("published"): + row["published"] = "" + rows.append( + [ome] + + [ + "" if row.get(c) is None else str(row.get(c, "")) + for c in self.columns + if c != "ome" + ] + ) + return rows + + def df2db( + self, + db_path: Optional[str] = None, + headers: bool = False, + paths: bool = False, + ) -> None: + """Write this MTDB as a tab-delimited `.mtdb` interchange file. + + With no `db_path` the database is printed to stdout, which is what makes + `mtdb extract | ...` composable. A file write is atomic.""" + rows = self._export_rows(paths=paths) + header = "#" + "\t".join(self.columns) if db_path: - with open(db_path, "w") as out: + with atomic_write(db_path) as out: if headers: - out.write("#" + "\t".join(self.columns) + "\n") - for ome in output: - if not paths: - for file_type in ["fna", "faa", "gff3"]: - output[ome][file_type] = output[ome][file_type].replace( - abb_paths[file_type][0] + ome + abb_paths[file_type][1], - "", - ) # abbreviate when possible - for rank in ["species", "genus", "strain"]: - try: - del output[ome]["taxonomy"][rank] - except (KeyError, TypeError) as e: - pass - - if output[ome]["taxonomy"]: - output[ome]["taxonomy"] = json.dumps(output[ome]["taxonomy"]) - else: - output[ome]["taxonomy"] = "{}" - if not output[ome]["published"]: - output[ome]["published"] = "" - out.write( - ome - + "\t" - + "\t".join([str(output[ome][x]) for x in output[ome]]) - + "\n" - ) + out.write(header + "\n") + for row in rows: + out.write("\t".join(row) + "\n") else: if headers: - print("#" + "\t".join(self.columns), flush=True) - - for ome in output: - if not paths: - for file_type in ["fna", "faa", "gff3"]: - output[ome][file_type] = output[ome][file_type].replace( - abb_paths[file_type][0] + ome + abb_paths[file_type][1], "" - ) # abbreviate when possible - for rank in ["species", "genus", "strain"]: - try: - del output[ome]["taxonomy"][rank] - except (KeyError, TypeError) as e: - pass - output[ome]["taxonomy"] = json.dumps(output[ome]["taxonomy"]) - print( - ome + "\t" + "\t".join([str(output[ome][x]) for x in output[ome]]), - flush=True, - ) + print(header, flush=True) + for row in rows: + print("\t".join(row), flush=True) + + def to_sql(self, db_path: str) -> str: + """Write this MTDB to a SQLite database, replacing it atomically.""" + return mtdb_sql.write_db(db_path, self.reset_index()) - def set_index(self, column="ome", inplace=False): + def set_index(self, column: Optional[str] = "ome", inplace: bool = False) -> "mtdb": data, retry, error, df, columns = ( {}, bool(column), @@ -247,7 +299,7 @@ def set_index(self, column="ome", inplace=False): ) if not column: return df.reset_index() - elif column not in {"assembly_acc", "ome"}: + elif column not in self._index_columns: raise KeyError(f'MTDB index must be "assembly_acc"/"ome"') elif df.index and not df.keys(): # empty df return mtdb({}, index=column) @@ -255,10 +307,9 @@ def set_index(self, column="ome", inplace=False): if not df["ome"]: return mtdb({}, index=column) while retry: - oldCol = set() try: columns.pop(columns.index(column)) - except (ValueError, IndexError) as e: + except (ValueError, IndexError): df = df.reset_index() # will this actually reset the index columns.pop(columns.index(column)) try: @@ -288,7 +339,7 @@ def set_index(self, column="ome", inplace=False): data[v][-1][head] = df[head][i] except KeyError: # if an index exists if error: - eprint("\nERROR: invalid column", flush=True) + logger.error("invalid column") return self df = df.reset_index() error = True @@ -297,7 +348,7 @@ def set_index(self, column="ome", inplace=False): retry = False return mtdb(data, column) - def reset_index(self): + def reset_index(self) -> "mtdb": df = copy.copy(self) if df.index: data = {x: [] for x in mtdb().columns} @@ -315,38 +366,336 @@ def reset_index(self): else: return df - def append(self, info={}): - # if any(x not in set(self.columns) for x in info): - # raise KeyError('Invalid keys: ' + str(set(info.keys()).difference(set(self.columns)))) + def append(self, info: Optional[Mapping[str, Any]] = None) -> "mtdb": + """Return a new MTDB with `info` added as a row. + + `copy.copy` is shallow, so the column lists have to be rebuilt rather + than appended to -- otherwise the returned MTDB shares its lists with + this one and appending mutates both.""" + if info is None: + info = {} index = self.index - df = copy.copy(self) - df = df.reset_index() + df = copy.copy(self).reset_index() info = { **info, **{k: None for k in set(self.columns).difference(set(info.keys()))}, } - for key in self.columns: - df[key].append(info[key]) + df = mtdb({key: list(df[key]) + [info[key]] for key in self.columns}) return df.set_index(index) + @staticmethod + def read_tax( + taxonomy_string: Union[str, Mapping[str, Any], None] + ) -> Dict[str, Any]: + """Read taxonomy from an MTDB by converting the string into a dictionary""" + tax_strs = [ + "superkingdom", + "kingdom", + "phylum", + "subphylum", + "class", + "order", + "family", + "subfamily", + ] + if taxonomy_string: + if isinstance(taxonomy_string, str): + dict_string = taxonomy_string.replace("'", '"') + try: + tax_dict = json.loads(dict_string) + except TypeError: + tax_dict = {} + else: + tax_dict = taxonomy_string + try: + tax_dict = { + **tax_dict, + **{x: "" for x in tax_strs if x not in tax_dict}, + } + except TypeError: # inappropriate tax_dict in the column + tax_dict = {x: "" for x in tax_strs} + return tax_dict + else: + return {} + + @staticmethod + def _reconcile_tax_dicts( + genera: Iterable[str], + tax_dicts: Mapping[str, Mapping[str, Any]], + forbid: Iterable[str], + ) -> Dict[str, Dict[str, Any]]: + """Drop empty/forbidden taxonomy entries and backfill missing genera""" + forbid = set(forbid) + tax_dicts = {x: tax_dicts[x] for x in tax_dicts if tax_dicts[x]} + for genus in tax_dicts: + tax_dicts[genus] = { + rank: name + for rank, name in tax_dicts[genus].items() + if rank not in forbid + } + for miss in set(genera).difference(tax_dicts.keys()): + tax_dicts[miss] = {} + return tax_dicts + + def prepare_tax_dicts( + self, tax_dicts: Optional[Dict[str, Any]] = None + ) -> "tuple[set, Dict[str, Any]]": + """Identify the genera that do not have higher taxonomy ascribed to them""" + if tax_dicts is None: + tax_dicts = {} + need_tax = set() + for ome, entry in self.set_index("ome").items(): + if entry["genus"] in tax_dicts: + continue + tax_json = self.read_tax(entry["taxonomy"]) + if any( + name + for rank, name in tax_json.items() + if rank not in {"genus", "species", "strain"} + ): + tax_dicts[entry["genus"]] = tax_json + else: + need_tax.add(entry["genus"]) + need_tax = need_tax.difference(tax_dicts.keys()) + return need_tax, tax_dicts + + def assimilate_tax( + self, + tax_dicts: Mapping[str, Mapping[str, Any]], + forbid: Optional[Iterable[str]] = None, + ) -> "tuple[mtdb, Dict[str, Any]]": + """Assign the resolved taxonomy dictionaries to this MTDB's taxonomy column""" + if forbid is None: + forbid = self._forbidden_tax_ranks + tax_dicts = self._reconcile_tax_dicts(set(self["genus"]), tax_dicts, forbid) + for i, genus in enumerate(self["genus"]): + self["taxonomy"][i] = tax_dicts[genus] + return tax_dicts + + def infer_rank(self, lineage: str) -> str: + """Identify the taxonomic rank associated with an inputted lineage of + interest""" + linlow = lineage.lower() + for ome, row in self.items(): + for rank, name in row["taxonomy"].items(): + if isinstance(name, str) and name.lower() == linlow: + return rank + + raise KeyError(f"no entry for {lineage}") + + def extract_unique( + self, allowed: int = 1, rank: str = "species", seed: Optional[int] = None + ) -> "mtdb": + """Extract unique rank from an MTDB. + + Genomes are sampled in a shuffled order, so `seed` is accepted to make a + selection reproducible.""" + keys = list(self.keys()) + random.Random(seed).shuffle(keys) + prep_db1 = mtdb().set_index("ome") + if rank == "strain": + found = set() + for ome in keys: + row = self[ome] + name = row["taxonomy"]["species"] + " " + row["strain"] + if name not in found: + prep_db1[ome] = row + found.add(name) + else: + found = defaultdict(int) + for ome in keys: + row = self[ome] + name = row["taxonomy"][rank] + found[name] += 1 + if found[name] <= allowed: + prep_db1[ome] = row + + return prep_db1 + + def extract_tax(self, lineages: Union[str, Iterable[str]]) -> "mtdb": + """Extract specific taxonomic lineages of interest based on their rank""" + if isinstance(lineages, str): + lineages = [lineages] + lineages = set(x.lower() for x in lineages) + rank_dict = {k: self.infer_rank(k) for k in list(lineages)} + ranks = list(set(rank_dict.values())) + + new_db = mtdb().set_index() + for ome in self: + for rank in ranks: + try: + if self[ome]["taxonomy"][rank].lower() in lineages: + new_db[ome] = self[ome] + except KeyError: # invalid rank key for row + pass # probably should standardize tax jsons period + + return new_db + + def extract_ome(self, omes: Iterable[str], column: str = "ome") -> "mtdb": + """Extract a list of genome codes (omes) of interest""" + omes = set(omes) # hoisted: rebuilding this per row is quadratic + new_db = mtdb().set_index(column) + db = self.set_index(column) + for i in db: + if i in omes: + new_db[i] = db[i] + return new_db.set_index() + + def extract_source(self, source: str) -> "mtdb": + """Extract an MTDB with genomes from a particular source""" + return mtdb( + { + ome: row + for ome, row in self.items() + if row["source"].lower() == source.lower() + }, + index="ome", + ) -def getLogin(ncbi, jgi): + def extract_pub(self) -> "mtdb": + """Extract only published and usable genomes""" + new_db = mtdb().set_index() + for ome, row in self.items(): + if row["published"]: + new_db[ome] = row + return new_db + + +def load_omes(db_path: str, omes: Iterable[str], add_paths: bool = True) -> "mtdb": + """Load only `omes` from the database at `db_path`. + + Against the SQLite backend this is an index seek per genome; against a + `.mtdb` flat file the whole file still has to be parsed, so the result is + the same either way and only the cost differs. This is the read that most + Mycotools commands actually want -- they operate on the genomes behind a + handful of accessions, not on the whole database.""" + db_path = format_path(db_path) + omes = set(omes) + if mtdb_sql.is_sqlite(db_path): + return mtdb(mtdb_sql.select_omes(db_path, omes, add_paths=add_paths)) + return mtdb(db_path, add_paths=add_paths).set_index("ome").extract_ome(omes) + + +def db_stem(db_path: Optional[str]) -> str: + """Basename of a database with its backend extension removed. + + Export filenames are built from this rather than from the primary's own + filename, so they stay `.mtdb` interchange files whichever backend they were + extracted from -- `20240101.mtdb` and `mtdb.db` give `20240101` and `mtdb`.""" + if not db_path: + return "mtdb" + name = Path(db_path).name + for suffix in (".mtdb", ".db"): + if name.endswith(suffix): + return name[: -len(suffix)] + return name + + +def omes_from_accessions(accs: Iterable[str]) -> "set[str]": + """MTDB aliases are `_`, so the ome is the leading field.""" + return {acc[: acc.find("_")] for acc in accs if "_" in acc} + + +# An NCBI api key is sent to the API as the `Api-Key` HTTP header, and the +# `datasets` CLI additionally echoes its own argv -- api key included -- into +# `X-Datasets-Client-Cmd`. Go's net/http refuses to transmit a header value +# holding a control character, so a key carrying the newline it was pasted with +# fails every request before it ever leaves the machine, reporting either +# `invalid header field value for "Api-Key"` or the same for +# `"X-Datasets-Client-Cmd"` depending on which header it validated first. +_ILLEGAL_HEADER_CHARS = re.compile(r"[\x00-\x1f\x7f]") + + +def clean_api_key(api_key): + """Remove characters that make an api key an illegal HTTP header value. + + Falsy keys pass through untouched so callers that distinguish `None` from + `""` keep doing so.""" + if not api_key: + return api_key + cleaned = _ILLEGAL_HEADER_CHARS.sub("", str(api_key)).strip() + if cleaned != api_key: + logger.warning( + "Removed whitespace/control characters from the NCBI api key; " + "they are not transmissible as an HTTP header" + ) + return cleaned - ncbi_email, ncbi_api, jgi_email, jgi_pwd = None, None, None, None + +def get_login(ncbi, jgi): + + ncbi_api, jgi_email, jgi_pwd = None, None, None print(flush=True) if ncbi: - ncbi_email = input("NCBI email: ") - ncbi_api = getpass.getpass(prompt="NCBI api key (blank if none): ") + ncbi_api = clean_api_key( + getpass.getpass(prompt="NCBI api key (blank if none): ") + ) if jgi: jgi_email = input("JGI email: ") jgi_pwd = getpass.getpass(prompt="JGI password (required): ") print(flush=True) - return ncbi_email, ncbi_api, jgi_email, jgi_pwd + return ncbi_api, jgi_email, jgi_pwd + + +# Path to the UNENCRYPTED credential store (see store_login). Kept separate from +# the password-encrypted key (`~/.mycotools/mtdb_key`) so the two never collide. +PLAIN_LOGIN_PATH = "~/.mycotools/mtdb_credentials.json" + + +def store_login( + ncbi_api, + jgi_email, + jgi_pwd, + info_path=PLAIN_LOGIN_PATH, + encrypted_path="~/.mycotools/mtdb_key", +): + """Store NCBI/JGI credentials WITHOUT a MycotoolsDB password. + + Credentials are written as JSON with owner-only (0600) permissions. Unlike + `encrypt_pw`, no password is set or required to read them back - trading + encryption for convenience. Anyone able to read the file can read the JGI + password in the clear, so the file permissions are its only protection. + """ + info_path = format_path(info_path) + Path(info_path).parent.mkdir(parents=True, exist_ok=True) + data = { + "ncbi_api": ncbi_api or "", + "jgi_email": jgi_email or "", + "jgi_pwd": jgi_pwd or "", + } + # create the file with restrictive permissions *before* writing the secret, + # so the password is never briefly exposed with a broader umask + fd = os.open(info_path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600) + with os.fdopen(fd, "w") as out: + json.dump(data, out) + os.chmod(info_path, 0o600) + logger.warning( + "Stored credentials UNENCRYPTED at %s (permissions 600). Anyone able to " + "read this file can read your JGI password.", + info_path, + ) + # a password-encrypted key would otherwise take precedence in login_check; + # remove it so the no-password store is the one that is actually used + enc = Path(format_path(encrypted_path)) + if enc.is_file(): + enc.unlink() + logger.info("Removed prior password-encrypted key %s", str(enc)) + + +def read_plain_login(info_path=PLAIN_LOGIN_PATH): + """Return (ncbi_api, jgi_email, jgi_pwd) from the unencrypted store written + by `store_login`. Missing fields come back as empty strings.""" + with open(format_path(info_path), "r") as raw: + data = json.load(raw) + return ( + clean_api_key(data.get("ncbi_api", "")), + data.get("jgi_email", ""), + data.get("jgi_pwd", ""), + ) def encrypt_pw( - ncbi_email, ncbi_api, jgi_email, jgi_pwd, @@ -356,7 +705,6 @@ def encrypt_pw( from cryptography.fernet import Fernet from cryptography.hazmat.primitives import hashes from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC - from cryptography.hazmat.backends import default_backend kdf = PBKDF2HMAC( algorithm=hashes.SHA256(), @@ -368,25 +716,32 @@ def encrypt_pw( hash_pwd, hash_check = True, False while hash_pwd != hash_check: if hash_pwd != True: - eprint("ERROR: passwords do not match", flush=True) + logger.error("passwords do not match") hash_pwd = getpass.getpass(prompt="New MycotoolsDB login password: ") hash_check = getpass.getpass(prompt="Confirm password: ") key = base64.urlsafe_b64encode(kdf.derive(hash_pwd.encode("utf-8"))) fernet = Fernet(key) - out_data = ncbi_email + "\t" + ncbi_api + "\t" + jgi_email + "\t" + jgi_pwd + out_data = ncbi_api + "\t" + jgi_email + "\t" + jgi_pwd encrypt_data = fernet.encrypt(out_data.encode("utf-8")) with open(format_path(info_path), "wb") as out: out.write(encrypt_data) + # keep credentials in exactly one place: drop any unencrypted store + plain = Path(format_path(PLAIN_LOGIN_PATH)) + if plain.is_file(): + plain.unlink() + logger.info("Removed unencrypted credential store %s", str(plain)) -def loginCheck(info_path="~/.mycotools/mtdb_key", ncbi=True, jgi=True, encrypt=False): +def login_check(info_path="~/.mycotools/mtdb_key", ncbi=True, jgi=True, encrypt=False): salt = b"D9\x82\xbfSibW(\xb1q\xeb\xd1\x84\x118" - # NEED to make this store a password - if os.path.isfile(format_path(info_path)): + # Credential source precedence: + # 1. password-encrypted key (encrypt_pw) - prompts for a password + # 2. unencrypted store (store_login) - no password required + # 3. interactive prompt (get_login) - not persisted + if Path(format_path(info_path)).is_file(): from cryptography.fernet import Fernet from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC - from cryptography.hazmat.backends import default_backend from cryptography.hazmat.primitives import hashes kdf = PBKDF2HMAC( @@ -402,27 +757,27 @@ def loginCheck(info_path="~/.mycotools/mtdb_key", ncbi=True, jgi=True, encrypt=F hash_pwd = sys.stdin.readline().rstrip() key = base64.urlsafe_b64encode(kdf.derive(hash_pwd.encode("utf-8"))) fernet = Fernet(key) - # with open(format_path(info_path) + '/.key', 'rb') as raw_key: - # fernet = Fernet(raw_key) with open(format_path(info_path), "rb") as raw_file: data = raw_file.read() decrypted = fernet.decrypt(data) - data = decrypted.decode("UTF-8").split() - if len(data) != 4: - eprint( - "BAD PASSWORD FILE. Delete ~/.mycotools/mtdb_key to reset.", flush=True - ) + data = decrypted.decode("UTF-8").split("\t") + # legacy key files stored a leading NCBI email; drop it if present + if len(data) == 4: + data = data[1:] + if len(data) != 3: + logger.error("BAD PASSWORD FILE. Delete ~/.mycotools/mtdb_key to reset.") sys.exit(8) - ncbi_email = data[0].rstrip() - ncbi_api = data[1].rstrip() - jgi_email = data[2].rstrip() - jgi_pwd = data[3].rstrip() + ncbi_api = clean_api_key(data[0]) + jgi_email = data[1].rstrip() + jgi_pwd = data[2].rstrip() + return ncbi_api, jgi_email, jgi_pwd + elif Path(format_path(PLAIN_LOGIN_PATH)).is_file(): + # unencrypted store written by store_login - no password required + return read_plain_login(PLAIN_LOGIN_PATH) else: - ncbi_email, ncbi_api, jgi_email, jgi_pwd = getLogin(ncbi, jgi) + ncbi_api, jgi_email, jgi_pwd = get_login(ncbi, jgi) # CURRENTLY THE REST DOESNT WORK, SO SKIP FOR NOW - return ncbi_email, ncbi_api, jgi_email, jgi_pwd - # - return ncbi_email, ncbi_api, jgi_email, jgi_pwd + return ncbi_api, jgi_email, jgi_pwd # opens a `log` file path to read, searches for the `ome` code followed by a whitespace character, and edits the line with `edit` @@ -445,7 +800,7 @@ def log_editor(log, ome, edit): towrite.write(new_data) -def readLog(log, columns="", sep="\t"): +def read_log(log, columns="", sep="\t"): log_dict = {} with open(log, "r") as raw: @@ -467,9 +822,13 @@ def readLog(log, columns="", sep="\t"): return log_dict -def primaryDB(path="$MYCODB", verbose=True): - """Acquire the path of the primary database by searching $MYCODB for a file - with a basename that starts with a date string %Y%m%d.""" +def primary_db(path="$MYCODB", verbose=True): + """Acquire the path of the primary database. + + A SQLite `mtdb.db` in $MYCODB is the primary database when present; + otherwise the newest dated `YYYYmmdd.mtdb` flat file is, which is what keeps + databases predating the SQLite backend usable. Callers only ever see a path + and pass it to `mtdb()`, which dispatches on the file's own format.""" path = path.replace("$", "") try: @@ -477,8 +836,12 @@ def primaryDB(path="$MYCODB", verbose=True): except KeyError: # $MYCODB not initialized return None + sql_path = format_path("$" + path + "/" + mtdb_sql.PRIMARY_DB_NAME) + if Path(sql_path).is_file(): + return sql_path + files = collect_files(full_path, "mtdb") - basenames = [os.path.basename(x) for x in files] + basenames = [Path(x).name for x in files] dates = [x.replace(".mtdb", "") for x in basenames if re.search(r"^\d+\.mtdb$", x)] primary = "19991231" # arbitrary primary for sorting for date in dates: @@ -488,7 +851,7 @@ def primaryDB(path="$MYCODB", verbose=True): primary = date if primary == "19991231": # if it is the arbitrary start if verbose: - eprint("\nWARNING: Primary MTDB not found in " + full_path, flush=True) + logger.warning("Primary MTDB not found in " + full_path) return None primary_path = format_path("$" + path + "/" + primary + ".mtdb") @@ -498,42 +861,25 @@ def primaryDB(path="$MYCODB", verbose=True): # imports database, converts into df # returns database dataframe def db2df(data, stdin=False): - """Deprecated legacy Pandas implementation of MTDB import""" - import pandas as pd, pandas + """Deprecated legacy Pandas implementation of MTDB import. + + Reading is delegated to the `mtdb` class so that a SQLite primary database, + a `.mtdb` interchange file, and an in-memory MTDB all behave identically + here; only the DataFrame conversion is still this function's own. The + previous implementation parsed the file a second time with `pd.read_csv` and + overwrote explicit `fna`/`faa`/`gff3` paths with $MYCO* defaults, silently + relocating standalone genomes.""" + import pandas as pd - columns = mtdb.columns if isinstance(data, mtdb): - db_df = pd.DataFrame(data.reset_index()) - elif not stdin: - data = format_path(data) - db_df = pd.read_csv(data, sep="\t") - if "ome" not in set(db_df.columns) and "assembly_acc" not in set(db_df.columns): - db_df = pd.read_csv(data, sep="\t", header=None) + db = data.reset_index() + elif stdin: + db = mtdb.from_string(data) else: - db_df = pd.read_csv(StringIO(data), sep="\t") - if "ome" not in set(db_df.columns) and "assembly_acc" not in set(db_df.columns): - db_df = pd.read_csv(StringIO(data), sep="\t", header=None) - - db_df = db_df.fillna("") - - db_df.columns = columns - for i, row in db_df.iterrows(): - db_df.at[i, "taxonomy"] = read_tax(row["taxonomy"]) - db_df.at[i, "taxonomy"]["genus"] = row["genus"] - db_df.at[i, "taxonomy"]["species"] = row["genus"] + " " + row["species"] - # if malformatted due to decreased entries in some lines, this will raise an IndexError - if ( - row["fna"] or row["fna"] == row["ome"] + ".fna" - ): # abbreviated line w/o file coordinates - db_df.at[i, "fna"] = os.environ["MYCOFNA"] + row["ome"] + ".fna" - db_df.at[i, "faa"] = os.environ["MYCOFAA"] + row["ome"] + ".faa" - db_df.at[i, "gff3"] = os.environ["MYCOGFF3"] + row["ome"] + ".gff3" - else: # has file coordinates - db_df.at[i, "fna"] = format_path(row["fna"]) - db_df.at[i, "faa"] = format_path(row["faa"]) - db_df.at[i, "gff3"] = format_path(row["gff3"]) - - return db_df + db = mtdb(format_path(data)).reset_index() + + db_df = pd.DataFrame({c: list(db[c]) for c in mtdb.columns}) + return db_df.fillna("") def df2std(df): @@ -550,7 +896,6 @@ def df2std(df): # if rescue is set to 0, do not output database if output dir does not exit def df2db(df, db_path, header=False, overwrite=False, std_col=True, rescue=True): """Deprecated output pandas MTDB implementation to file""" - import pandas as pd, pandas df = df.set_index("ome") df = df.sort_index() @@ -560,9 +905,9 @@ def df2db(df, db_path, header=False, overwrite=False, std_col=True, rescue=True) db_path = format_path(db_path) elif overwrite: number = 0 - while os.path.exists(db_path): + while Path(db_path).exists(): number += 1 - db_path = os.path.normpath(db_path) + "_" + str(number) + db_path = str(Path(db_path)) + "_" + str(number) if std_col: df = df2std(df) @@ -573,24 +918,21 @@ def df2db(df, db_path, header=False, overwrite=False, std_col=True, rescue=True) break except FileNotFoundError: if rescue: - eprint( - "\nOutput directory does not exist. Attempting to save in home folder.", - flush=True, + logger.warning( + "Output directory does not exist. Attempting to save in home folder." ) - db_path = "~/" + os.path.basename(os.path.normpath(db_path)) + db_path = "~/" + Path(str(Path(db_path))).name df.to_csv(db_path, sep="\t", index=None) raise FileNotFoundError break else: - eprint( - "\nOutput directory does not exist. Rescue not enabled.", flush=True - ) + logger.error("Output directory does not exist. Rescue not enabled.") raise FileNotFoundError break def hit2taxonomy( - taxid, rank="kingdom", lineage="fungi", skip=False, email=None, api=None + taxid, rank="kingdom", lineage="fungi", skip=False, api=None ): """ Takes a searchTerm string, queries NCBI via Entrez, obtains TaxIDs, @@ -608,7 +950,7 @@ def hit2taxonomy( except urllib.error.HTTPError as goon: count += 1 if count == 5 and skip: - print("\n5 failed HTTP queries. Is NCBI down?", flush=True) + logger.error("5 failed HTTP queries. Is NCBI down?") sys.exit(100) if 500 <= goon.code <= 599: time.sleep(1) @@ -623,10 +965,10 @@ def hit2taxonomy( time.sleep(1) count += 1 if count == 5: - eprint("\nERROR: 5 failed taxonomy queries. " + str(taxid), flush=True) - eprint(tax_handle, flush=True) + logger.error("5 failed taxonomy queries. " + str(taxid)) + logger.debug("%s", tax_handle) for line in tax_handle: - print(line, flush=True) + logger.debug(line) if skip: sys.exit(1) records = False @@ -635,11 +977,10 @@ def hit2taxonomy( if "latin-1" in str(tax_handle): count = 0 - eprint("\tERROR: latin-1 encoding", flush=True) + logger.error("latin-1 encoding") for line in tax_handle: - print(line, flush=True) + logger.debug(line) time.sleep(30) - Entrez.email = email Entrez.api_key = api sleep = True @@ -656,40 +997,34 @@ def hit2taxonomy( return tax_dict, sleep -def prepare_tax_dicts(df, tax_dicts={}): - """Identify the genera that do not have higher taxonomy ascribed to them""" +def prepare_tax_dicts(df, tax_dicts=None): + """Identify the genera that do not have higher taxonomy ascribed to them. - need_tax = set() + Backwards-compatible dispatcher: MTDBs (and MTDB-shaped dicts) use + ``mtdb.prepare_tax_dicts``; a deprecated pandas DataFrame uses the path + below.""" + if tax_dicts is None: + tax_dicts = {} # is the df of mtdb class or have a taxonomy column as a list? - if isinstance(df, mtdb) or isinstance(df["taxonomy"], list): - df = df.set_index("ome") - for k, v in df.items(): - if v["genus"] in tax_dicts: - continue - tax_json = read_tax(v["taxonomy"]) - if any( - v - for k, v in tax_json.items() - if k not in {"genus", "species", "strain"} - ): - tax_dicts[v["genus"]] = tax_json - else: - need_tax.add(v["genus"]) + if isinstance(df, mtdb): + return df.prepare_tax_dicts(tax_dicts) + if isinstance(df["taxonomy"], list): + return mtdb(df).prepare_tax_dicts(tax_dicts) # otherwise it is a pandas dataframe - else: - df["taxonomy"] = df["taxonomy"].fillna({}) - for k, v in df.iterrows(): - if v["genus"] in tax_dicts: - continue - tax_json = read_tax(v["taxonomy"]) - if any( - v - for k, v in tax_json.items() - if k not in {"genus", "species", "strain"} - ): - tax_dicts[v["genus"]] = tax_json - else: - need_tax.add(v["genus"]) + need_tax = set() + df["taxonomy"] = df["taxonomy"].fillna({}) + for k, v in df.iterrows(): + if v["genus"] in tax_dicts: + continue + tax_json = read_tax(v["taxonomy"]) + if any( + name + for rank, name in tax_json.items() + if rank not in {"genus", "species", "strain"} + ): + tax_dicts[v["genus"]] = tax_json + else: + need_tax.add(v["genus"]) need_tax = set(need_tax).difference(set(tax_dicts.keys())) return need_tax, tax_dicts @@ -712,7 +1047,7 @@ def query_ncbi4taxonomy(genus, api_key, king, rank, count=0): count += 1 if not ids: - eprint(f"\t\t{genus} TaxID acquisition failed", flush=True) + logger.warning(f"{genus} TaxID acquisition failed") return None, count # for each taxID acquired, fetch the actual taxonomy information @@ -752,14 +1087,14 @@ def query_ncbi4taxonomy(genus, api_key, king, rank, count=0): if lineage["ScientificName"].lower() == king.lower(): taxid = tax if len(ids) > 1: - print("\t\tMultiple Tax IDs: " + str(ids), flush=True) + logger.warning("Multiple Tax IDs: " + str(ids)) break else: for lineage in lineages: taxid = tax if taxid == 0: - eprint(f"\t\t{genus} not recovered in {king}", flush=True) + logger.warning(f"{genus} not recovered in {king}") return None, count # for each taxonomic classification, add it to the taxonomy dictionary string @@ -797,7 +1132,7 @@ def gather_taxonomy( with open(output_path + ".tmp", "w") as out: for genus, tax_dict in tax_dicts.items(): out.write(f"{genus}\t{json.dumps(tax_dict)}\n") - os.rename(f"{output_path}.tmp", output_path) + Path(f"{output_path}.tmp").rename(output_path) return tax_dicts @@ -824,7 +1159,7 @@ def gather_taxonomy_dataset( with open(tax_accs_file, "w") as out: out.write("\n".join(sorted(need_tax))) - cwd = os.getcwd() + cwd = str(Path.cwd()) os.chdir(output_path) cmd_scaf = [ "datasets", @@ -834,6 +1169,7 @@ def gather_taxonomy_dataset( "--inputfile", tax_accs_file, ] + api_key = clean_api_key(api_key) if api_key: cmd_scaf.extend(["--api-key", api_key]) @@ -845,19 +1181,19 @@ def gather_taxonomy_dataset( count = 0 dataset_path = output_path + "ncbi_dataset.zip" while count < 3: - cmd_call = subprocess.call(cmd_scaf, stdout=v, stderr=v) + subprocess.call(cmd_scaf, stdout=v, stderr=v) try: with zipfile.ZipFile(dataset_path, "r") as zip_ref: zip_ref.extractall(zip_ref) except zipfile.BadZipFile: count += 1 if count == 3: - eprint("ERROR: taxonomy acquisition failed", flush=True) + logger.error("taxonomy acquisition failed") sys.exit(130) continue break - os.remove(dataset_path) + Path(dataset_path).unlink() unzip_path = output_path + "ncbi_dataset/" tax_dicts = parse_dataset_taxonomy( @@ -901,81 +1237,18 @@ def parse_dataset_taxonomy(tax_json, tax_dicts, tax_head, rank_head): return tax_dicts -def read_tax(taxonomy_string): - """Read taxonomy from an MTDB by converting the string into a dictionary""" - tax_strs = [ - "superkingdom", - "kingdom", - "phylum", - "subphylum", - "class", - "order", - "family", - "subfamily", - ] - if taxonomy_string: - if isinstance(taxonomy_string, str): - dict_string = taxonomy_string.replace("'", '"') - try: - tax_dict = json.loads(dict_string) - except TypeError: - tax_dict = {} - else: - tax_dict = taxonomy_string - try: - tax_dict = {**tax_dict, **{x: "" for x in tax_strs if x not in tax_dict}} - except TypeError: # inappropriate tax_dict in the column - tax_dict = {x: "" for x in tax_strs} - return tax_dict - else: - return {} - - -# assimilate taxonomy dictionary strings and append the resulting taxonomy string dicts to an inputted database -# forbid a list of taxonomic classifications you are not interested in and return a new database -def assimilate_tax( - db, - tax_dicts, - ome_index="ome", - forbid={ - "no rank", - "subkingdom", - "genus", - "species", - "species group", - "varietas", - "forma", - }, -): - - genera = set(db["genus"]) - tax_dicts = {x: tax_dicts[x] for x in tax_dicts if tax_dicts[x]} - for genus in tax_dicts: - tax_dicts[genus] = { - x: tax_dicts[genus][x] for x in tax_dicts[genus] if x not in forbid - } - missing = list(genera.difference(set(tax_dicts.keys()))) - - for miss in missing: - tax_dicts[miss] = {} - if isinstance(db, mtdb): - for i, genus in enumerate(db["genus"]): - db["taxonomy"][i] = tax_dicts[genus] - return mtdb(db), tax_dicts - else: - for i, row in db.iterrows(): - db.at[i, "taxonomy"] = tax_dicts[row["genus"]] - - return db, tax_dicts +# read_tax is defined as mtdb.read_tax; expose it at module scope for the +# historical ``from mycotools.lib.dbtools import read_tax`` import. +read_tax = mtdb.read_tax def parse_user_config(mtdb_config_file=format_path("~/.mycotools/config.json")): config_dir = format_path("~/.mycotools/") - if not os.path.isdir(config_dir): - os.mkdir(config_dir) + if not Path(config_dir).is_dir(): + Path(config_dir).mkdir() config_dir += "/" - if os.path.isfile(mtdb_config_file): + if Path(mtdb_config_file).is_file(): config = read_json(mtdb_config_file) else: config = {"log": {}} @@ -1008,9 +1281,9 @@ def mtdb_initialize( mtdb_config = read_json(mycodb_loc + "config/mtdb.json") dbtype = mtdb_config["branch"] - eprint("Establishing " + dbtype + " connection", flush=True) + logger.info("Establishing " + dbtype + " connection") - if not os.path.isdir(mycodb_loc + "mtdb/") and not init: + if not Path(mycodb_loc + "mtdb/").is_dir() and not init: raise FileNotFoundError("invalid MycotoolsDB path") dPath = mycodb_loc + "data/" user_config[mycodb_loc] = { @@ -1021,7 +1294,6 @@ def mtdb_initialize( } login_time = datetime.datetime.now().strftime("%Y%m%d %H:%M:%S") - # login_time = datetime.datetime.now().strftime('%Y%m%d') if dbtype in user_config["log"]: user_config["log"][dbtype][mycodb_loc] = login_time else: @@ -1031,11 +1303,11 @@ def mtdb_initialize( interface = format_path("~/.mycotools/config.json") -if os.path.isfile(interface): +if Path(interface).is_file(): envs_info = read_json(interface) if envs_info["active"]: for var, env in envs_info[envs_info["active"]].items(): os.environ[var] = env -# if not primaryDB(): +# if not primary_db(): # eprint('WARNING: Primary MycotoolsDB not connected; setup using `mtdb u/-i/-p/-f`', flush = True) diff --git a/mycotools/lib/kontools.py b/mycotools/lib/kontools.py index 4421b3e..af86c5d 100755 --- a/mycotools/lib/kontools.py +++ b/mycotools/lib/kontools.py @@ -6,15 +6,64 @@ import glob import gzip import json +import logging import shutil import tarfile import argparse import subprocess +from pathlib import Path +from contextlib import contextmanager from tqdm import tqdm from datetime import datetime -class kon_log: +logger = logging.getLogger(__name__) + + +class _LevelFormatter(logging.Formatter): + """Prefix records with their level name, except INFO, which is emitted + verbatim so status messages read as they did before the logging + migration.""" + + _PREFIX = { + logging.DEBUG: "DEBUG: ", + logging.INFO: "", + logging.WARNING: "WARNING: ", + logging.ERROR: "ERROR: ", + logging.CRITICAL: "CRITICAL: ", + } + + def format(self, record): + message = super().format(record) + prefix = self._PREFIX.get(record.levelno, "") + # Do not double-prefix messages that already carry their level word. + if prefix and message.lstrip().upper().startswith(prefix.strip().rstrip(":")): + prefix = "" + return prefix + message + + +def setup_logging(verbose=False, level=None): + """Configure logging for the Mycotools CLI. + + Attaches a single stderr handler to the root logger (idempotent across + repeated calls) and sets the ``mycotools`` logger to INFO, or DEBUG when + ``verbose`` is True. The root logger stays at WARNING so third-party + libraries do not clutter diagnostic output. Call once from a script's + entry point after parsing arguments.""" + if level is None: + level = logging.DEBUG if verbose else logging.INFO + root = logging.getLogger() + root.setLevel(logging.WARNING) + if not any(getattr(h, "_mycotools", False) for h in root.handlers): + handler = logging.StreamHandler(sys.stderr) +# handler.setFormatter(_LevelFormatter("%(message)s")) + handler._mycotools = True + root.addHandler(handler) + logging.getLogger("mycotools").setLevel(level) + return logging.getLogger("mycotools") + + +class KonLog: """A print class designed to enable swift string formatting while moving between scripts""" @@ -127,7 +176,7 @@ def parse_run_log(log_path, args_dict, fail=set()): if isinstance(args_dict, argparse.Namespace): args_dict = namespace_to_dict(args_dict) - if not os.path.isfile(log_path): + if not Path(log_path).is_file(): write_json(args_dict, log_path) return {} else: @@ -151,7 +200,7 @@ def hex2rgb(hexCode): return tuple(int(hexCode.lstrip("#")[i : i + 2], 16) for i in (0, 2, 4)) -def getColors(size, ignore=[], rgb=False): +def get_colors(size, ignore=[], rgb=False): if size < 16: colors = [ "#000000", @@ -286,28 +335,28 @@ def getColors(size, ignore=[], rgb=False): def tardir(dir_, rm=True): - if not os.path.isdir(format_path(dir_)): + if not Path(format_path(dir_)).is_dir(): return False with tarfile.open(format_path(dir_)[:-1] + ".tar.gz", "w:gz") as tar: - tar.add(dir_, arcname=os.path.basename(format_path(dir_)[:-1])) + tar.add(dir_, arcname=Path(format_path(dir_)[:-1]).name) if rm: shutil.rmtree(dir_) def untardir(dir_, rm=False, to=None): if not to: - to = os.path.dirname(dir_[:-1]) + to = str(Path(dir_[:-1]).parent) tar = tarfile.TarFile.open(dir_) tar.extractall(path=to) tar.close() if rm: - os.remove(dir_) + Path(dir_).unlink() def checkdir(dir_, unzip=False, to=None, rm=False): - if os.path.isdir(dir_): + if Path(dir_).is_dir(): return True - elif os.path.isfile(format_path(dir_) + ".tar.gz"): + elif Path(format_path(dir_) + ".tar.gz").is_file(): if unzip: if dir_.endswith("/"): dir_ = dir_[:-1] @@ -316,31 +365,6 @@ def checkdir(dir_, unzip=False, to=None, rm=False): return False -def eprint(*args, **kwargs): - """Prints to stderr""" - print(*args, file=sys.stderr, **kwargs) - - -def fprint(out_str, log): - with open(log, "a") as out: - out.write(args) - - -def zprint(out_str, log=None, flush=True): - fprint(out_str, log) - print(out_str, flush=flush) - - -def vprint(toPrint, v=False, e=False, flush=True): - """Boolean print option to stdout or stderr (e)""" - - if v: - if e: - eprint(toPrint, flush=True) - else: - print(toPrint, flush=True) - - def read_json(config_path, compress=False): if compress or config_path.endswith(".gz"): @@ -362,6 +386,32 @@ def write_json(obj, json_path, compress=False, indent=1, **kwargs): json.dump(obj, json_out, indent=indent, **kwargs) +@contextmanager +def atomic_write(path, mode="w", suffix=".tmp", **kwargs): + """Context manager for crash-safe file writes. + + Yields a handle to a temporary sibling file (`path` + `suffix`) and, only + upon leaving the block without error, atomically replaces `path` with it. + Because the temporary file lives beside the target it shares a filesystem, + so the final replace is atomic and can never leave a half-written `path` + behind. If the block raises, the temporary file is removed and `path` is + left untouched. + + Consolidates the `open(path + '.tmp')` ... `shutil.move`/`Path.rename` + idiom repeated throughout the codebase. + """ + path = Path(path) + tmp_path = path.with_name(path.name + suffix) + try: + with open(tmp_path, mode, **kwargs) as handle: + yield handle + tmp_path.replace(path) + except BaseException: + if tmp_path.is_file(): + tmp_path.unlink() + raise + + def gunzip(gzip_file, remove=True, spacer="\t"): """gunzips gzip_file and removes if successful""" @@ -372,12 +422,12 @@ def gunzip(gzip_file, remove=True, spacer="\t"): for line in f_in: f_out.write(line) if remove: - os.remove(gzip_file) + Path(gzip_file).unlink() return new_file except: - if os.path.isfile(new_file): - if os.path.isfile(gzip_file): - os.remove(new_file) + if Path(new_file).is_file(): + if Path(gzip_file).is_file(): + Path(new_file).unlink() raise IOError("gunzip " + str(gzip_file) + " failed") @@ -466,96 +516,77 @@ def fmt_float(val, sig_dig=None): return val_str -def findExecs(deps, exit=set(), verbose=True): +def find_execs(deps, exit=set(), verbose=True): """ Inputs list of dependencies, `dep`, to check path. If dependency is in exit and dependency is not in path, then exit. """ - vprint("\nDependency check:", v=verbose, e=True, flush=True) + logger.debug("Dependency check:") checks, failed = [], [] if isinstance(deps, str): deps = [deps] for dep in sorted(deps): check = shutil.which(dep) - vprint("{:<15}".format(dep + ":", flush=True) + str(check), v=verbose, e=True) + logger.debug("{:<15}".format(dep + ":") + str(check)) if not check and dep in exit: failed.append(dep) else: checks.append(check) if failed: - eprint("\nERROR: missing dependencies:", flush=True) - for f in failed: - vprint(f, v=verbose, e=True) - eprint() + logger.error("missing dependencies: " + ", ".join(failed)) sys.exit(300) return checks -def findEnvs(envs, exit=set(), verbose=True): +def find_envs(envs, exit=set(), verbose=True): """ Inputs list of paths, `envs`, to check path. If env is not in path and it is in exit, exit. """ - vprint("\nEnvironment check:", v=verbose, e=True, flush=True) + logger.debug("Environment check:") if type(envs) is str: envs = [envs] - eprint(flush=True) for env in envs: try: - vprint( - "{:<15}".format(env + ":", flush=True) + str(os.environ[env]), - v=verbose, - e=True, - ) + logger.debug("{:<15}".format(env + ":") + str(os.environ[env])) except KeyError: - vprint("{:<15}".format(env + ":", flush=True) + "None", v=verbose, e=True) + logger.debug("{:<15}".format(env + ":") + "None") if env in exit: - eprint("\nERROR: " + env + " not in PATH", flush=True) + logger.error(env + " not in PATH") sys.exit(301) -def expandEnvVar(path): - """Expands environment variables by regex substitution""" - - envs = re.findall(r"\$[^/]+", path) - for env in envs: - path = path.replace(env, os.environ[env.replace("$", "")]) - - return path.replace("//", "/") - - def format_path(path, force_dir=False): - """Goal is to convert all path types to absolute path with explicit dirs""" + """Convert a path to an absolute path with explicit directories. + + Expands ``~`` and ``$VAR`` (raising KeyError on an undefined variable). By + convention an existing directory is returned with a trailing ``/`` while + files and non-existent paths have none; ``force_dir`` forces a trailing + ``/`` for directories that do not exist yet. Symlinks are not resolved.""" - # path = path.replace('//','/') - # except AttributeError: # not a string - # return None # removed this because let it be handled on the other end - # try: if path: - path = os.path.expanduser(path) - path = expandEnvVar(path) - # path = os.path.abspath( path ) - # except TypeError: - # return None again, let this be handled on the other end to increase - # throughput + path = str(Path(str(path)).expanduser()) + for env in re.findall(r"\$[^/]+", path): + path = path.replace(env, os.environ[env.replace("$", "")]) + path = path.replace("//", "/") if force_dir: if not path.endswith("/"): path += "/" else: if path.endswith("/"): - if not os.path.isdir(path): + if not Path(path).is_dir(): path = path[:-1] else: - if os.path.isdir(path): + if Path(path).is_dir(): path += "/" if not path.startswith("/"): - path = os.getcwd() + "/" + path + path = str(Path.cwd()) + "/" + path path = path.replace("/./", "/") while "/../" in path: @@ -580,17 +611,12 @@ def collect_files(directory="./", filetype="*", recursive=False): else: filetypes = [filetype] - directory = format_path(directory) + directory = Path(format_path(directory)) filelist = [] for filetype in filetypes: - if recursive: - filelist.extend( - glob.glob(directory + "/**/*." + filetype, recursive=recursive) - ) - else: - filelist.extend( - glob.glob(directory + "/*." + filetype, recursive=recursive) - ) + pattern = "*." + filetype + matches = directory.rglob(pattern) if recursive else directory.glob(pattern) + filelist.extend(str(match) for match in matches) return filelist @@ -603,10 +629,10 @@ def collect_dirs(input_glob, recursive=False): """ in_dirs = glob.glob(input_glob, recursive=recursive) - return [dir_ for dir_ in in_dirs if os.path.isdir(dir_)] + return [dir_ for dir_ in in_dirs if Path(dir_).is_dir()] -def dictSplit(Dict, factor): +def dict_split(Dict, factor): """ Inputs: a dictionary `Dict`, and an integer `factor` to split by Outputs: a list of split dictionaries `list_dict` @@ -647,13 +673,13 @@ def sys_start(args, usage, min_len, dirs=[], files=[]): elif len(args) < min_len: print("\n" + usage + "\n", flush=True) sys.exit(1) - elif not all(os.path.isfile(format_path(x)) for x in files): + elif not all(Path(format_path(x)).is_file() for x in files): print("\n" + usage, flush=True) - eprint("ERROR: input file(s) do not exist\n", flush=True) + logger.error("input file(s) do not exist") sys.exit(3) - elif not all(os.path.isfile(format_path(x)) for x in dirs): + elif not all(Path(format_path(x)).is_file() for x in dirs): print("\n" + usage, flush=True) - eprint("ERROR: input directory does not exist\n", flush=True) + logger.error("input directory does not exist") sys.exit(4) return args @@ -663,7 +689,6 @@ def inject_args(args, injection_calls): manual_cmds = [] for in_call in injection_calls: if in_call in args: - prohibited = {";", "&", "&&", "\n", "\r"} man_index = args.index(in_call) for char in args[man_index + 1]: if char == '"' or char == "'": @@ -700,7 +725,6 @@ def intro(script_name, args_dict, credit="", log=False, stdout=True): """ start_time = datetime.now() - date = start_time.strftime("%Y%m%d") out_str = ( "\n" + script_name + "\n" + credit + "\nExecution began: " + str(start_time) @@ -709,12 +733,7 @@ def intro(script_name, args_dict, credit="", log=False, stdout=True): for arg in args_dict: out_str += "\n" + "{:<30}".format(arg.upper() + ":") + str(args_dict[arg]) - if log: - zprint(out_str, log) - elif stdout: - print(out_str, flush=True) - else: - eprint(out_str, flush=True) + logger.info(out_str) return start_time @@ -736,12 +755,7 @@ def outro(start_time, log=False, stdout=True): + " minutes\n" ) - if log: - zprint(out_str, log) - elif not stdout: - eprint(out_str, flush=True) - else: - print(out_str, flush=True) + logger.info(out_str) sys.exit(0) @@ -754,17 +768,17 @@ def prep_output(output, mkdir=True, require_newdir=False, cd=False): """ output = format_path(output, force_dir=True) - if os.path.isdir(output): + if Path(output).is_dir(): if require_newdir: - eprint("\nERROR: directory exists.", flush=True) + logger.error("directory exists.") return None - elif os.path.exists(output): - output = os.path.dirname(output) + elif Path(output).exists(): + output = str(Path(output).parent) else: if not mkdir: - eprint("\nERROR: directory does not exist.", flush=True) + logger.error("directory does not exist.") return None - os.mkdir(output) + Path(output).mkdir() if cd: os.chdir(output) @@ -772,10 +786,10 @@ def prep_output(output, mkdir=True, require_newdir=False, cd=False): return output -def mkOutput(base_dir, program, reuse=True, suffix=datetime.now().strftime("%Y%m%d")): +def mk_output(base_dir, program, reuse=True, suffix=datetime.now().strftime("%Y%m%d")): if not base_dir: - base_dir = os.getcwd() + "/" - if not os.path.isdir(format_path(base_dir)): + base_dir = str(Path.cwd()) + "/" + if not Path(format_path(base_dir)).is_dir(): raise FileNotFoundError(base_dir + " does not exist") if suffix: out_dir = format_path(base_dir) + program + "_" + suffix @@ -784,18 +798,18 @@ def mkOutput(base_dir, program, reuse=True, suffix=datetime.now().strftime("%Y%m if not reuse: count, count_dir = 1, out_dir - while os.path.isdir(count_dir): + while Path(count_dir).is_dir(): count_dir += "_" + str(count) count += 1 - os.mkdir(count_dir) + Path(count_dir).mkdir() return count_dir + "/" else: - if not os.path.isdir(out_dir): - os.mkdir(out_dir) + if not Path(out_dir).is_dir(): + Path(out_dir).mkdir() return out_dir + "/" -def checkDep(dep_list=[], var_list=[], exempt=set()): +def check_dep(dep_list=[], var_list=[], exempt=set()): """Checks all dependencies in path from list, optional exemption set""" failedVars = [] @@ -809,12 +823,12 @@ def checkDep(dep_list=[], var_list=[], exempt=set()): check.append(False) failedVars.append(var) if not all(check): - eprint("\nERROR: Dependencies not met:", flush=True) + logger.error("Dependencies not met:") for dep in dep_list: if not shutil.which(dep): - eprint(dep + " not in PATH", flush=True) + logger.error(dep + " not in PATH") for failed in failedVars: - eprint(failed + " variable not set", flush=True) + logger.error(failed + " variable not set") sys.exit(135) diff --git a/mycotools/lib/mtdb_sql.py b/mycotools/lib/mtdb_sql.py new file mode 100644 index 0000000..f7e15b0 --- /dev/null +++ b/mycotools/lib/mtdb_sql.py @@ -0,0 +1,604 @@ +#! /usr/bin/env python3 +"""SQLite storage backend for MycotoolsDB (MTDB). + +The primary MTDB is stored as a single SQLite file (``$MYCODB/mtdb.db``); the +tab-delimited ``.mtdb`` file remains the interchange format, read and written by +``mycotools.lib.dbtools.mtdb``. Nothing here is user-facing: the ``mtdb`` class +is still the interface, and every call site that passes a path to ``mtdb()`` +keeps working because :func:`is_sqlite` dispatches on the file itself. + +Two properties motivate the backend: + +* **Indexed lookup.** ``mtdb ``, ``mtdb accession``, and ``mtdb files`` + answer questions about a handful of genomes. Against a flat file each has to + parse the whole database; here they are index seeks (:func:`select_omes`). +* **Normalized taxonomy.** A lineage belongs to a genus, not to a genome, so it + is stored once in the ``taxonomy`` table instead of being repeated as JSON on + every row -- ~70% of a flat ``.mtdb`` is duplicated lineage text. + +Writes go through :func:`write_db`, which builds into a temporary file and +renames over the target, so a cancelled update can never leave a partial +database where ``primary_db()`` will find it. + +Path columns follow the flat-file convention: an empty ``fna``/``faa``/``gff3`` +means "the default ``$MYCOFNA``/``$MYCOFAA``/``$MYCOGFF3`` location for this +ome". Storing them abbreviated keeps a database valid after its root moves; +:func:`rows2columns` expands them on read. +""" + +from __future__ import annotations + +import json +import logging +import os +import sqlite3 +from pathlib import Path +from typing import Any, Dict, Iterable, List, Mapping, Optional, Sequence + +logger = logging.getLogger(__name__) + +#: bumped when the on-disk layout changes in a way older code cannot read +SCHEMA_VERSION = 1 + +#: basename of the SQLite primary database within $MYCODB +PRIMARY_DB_NAME = "mtdb.db" + +#: first 16 bytes of any SQLite 3 file +_SQLITE_MAGIC = b"SQLite format 3\x00" + +#: MTDB columns as they appear in a `.mtdb` row, in order. Mirrors +#: ``mtdb.columns``; duplicated here so this module does not import dbtools +#: (dbtools imports *it*). +COLUMNS = [ + "ome", + "genus", + "species", + "strain", + "taxonomy", + "version", + "source", + "biosample", + "assembly_acc", + "acquisition_date", + "published", + "fna", + "faa", + "gff3", +] + +#: columns physically stored on the `genome` table -- `taxonomy` is normalized +#: out into its own table and reattached on read +_GENOME_COLUMNS = [c for c in COLUMNS if c != "taxonomy"] + +#: file-path columns and the environment variable holding their default dir +_PATH_COLUMNS = {"fna": "MYCOFNA", "faa": "MYCOFAA", "gff3": "MYCOGFF3"} + +#: ranks backfilled empty when a lineage omits them. Must stay identical to +#: ``mtdb.read_tax``'s list -- the two readers have to produce the same taxonomy +#: dict for the same genome, or a flat/SQLite round trip is not an identity. +#: Any other rank present in the source (``clade``, ``subclass``, ...) is stored +#: and returned as-is; this list only governs what gets added. +_TAX_RANKS = ( + "superkingdom", + "kingdom", + "phylum", + "subphylum", + "class", + "order", + "family", + "subfamily", +) + +_SCHEMA = """ +CREATE TABLE IF NOT EXISTS genome ( + ome TEXT PRIMARY KEY, + genus TEXT NOT NULL DEFAULT '', + species TEXT NOT NULL DEFAULT '', + strain TEXT NOT NULL DEFAULT '', + version TEXT NOT NULL DEFAULT '', + source TEXT NOT NULL DEFAULT '', + biosample TEXT NOT NULL DEFAULT '', + assembly_acc TEXT NOT NULL DEFAULT '', + acquisition_date TEXT NOT NULL DEFAULT '', + published TEXT NOT NULL DEFAULT '', + fna TEXT NOT NULL DEFAULT '', + faa TEXT NOT NULL DEFAULT '', + gff3 TEXT NOT NULL DEFAULT '' +); + +CREATE TABLE IF NOT EXISTS taxonomy ( + genus TEXT PRIMARY KEY, + lineage TEXT NOT NULL DEFAULT '{}' +); + +CREATE INDEX IF NOT EXISTS genome_assembly_acc ON genome(assembly_acc); +CREATE INDEX IF NOT EXISTS genome_genus ON genome(genus); +CREATE INDEX IF NOT EXISTS genome_source ON genome(source); +CREATE INDEX IF NOT EXISTS genome_published ON genome(published); +""" + + +class MTDBSchemaError(Exception): + """Raised when a SQLite MTDB was written by an incompatible version.""" + + +# --------------------------------------------------------------------------- # +# detection / connection +# --------------------------------------------------------------------------- # +def is_sqlite(path: "str | Path | None") -> bool: + """Is `path` a SQLite file? Detects by magic bytes, not by extension. + + Every entry point that accepts a database path funnels through here, which + is what lets a `.mtdb` flat file and a SQLite database be used + interchangeably wherever a path is currently accepted.""" + if not path: + return False + try: + with open(path, "rb") as raw: + return raw.read(16) == _SQLITE_MAGIC + except (OSError, TypeError, ValueError): + return False + + +def connect(path: "str | Path", read_only: bool = True) -> sqlite3.Connection: + """Open a SQLite MTDB and verify its schema version. + + Read-only connections use a URI so concurrent readers can never be blocked + by, or interfere with, a writer.""" + path = str(path) + if read_only: + conn = sqlite3.connect(f"file:{path}?mode=ro", uri=True) + else: + conn = sqlite3.connect(path) + conn.row_factory = sqlite3.Row + version = conn.execute("PRAGMA user_version").fetchone()[0] + if version > SCHEMA_VERSION: + conn.close() + raise MTDBSchemaError( + f"{path} uses MTDB schema v{version}; this Mycotools reads " + f"v{SCHEMA_VERSION}. Upgrade Mycotools." + ) + return conn + + +def _init_schema(conn: sqlite3.Connection) -> None: + conn.executescript(_SCHEMA) + conn.execute(f"PRAGMA user_version = {SCHEMA_VERSION}") + + +# --------------------------------------------------------------------------- # +# row <-> column-dict conversion +# --------------------------------------------------------------------------- # +def _resolve_paths(row: Dict[str, Any], ome: str, env: Optional[Mapping[str, str]]) -> None: + """Expand abbreviated path columns in place, as the flat reader does. + + `env` is the pre-read {var: dir} mapping; None means the process is not + linked to a primary MTDB, in which case abbreviated paths stay empty and the + caller is responsible for reporting it.""" + if env is None: + return + for col, var in _PATH_COLUMNS.items(): + value = row.get(col) or "" + if not value or value == f"{ome}.{col}": + row[col] = env[var] + ome + "." + col + + +def read_path_env() -> Optional[Dict[str, str]]: + """Read the three MTDB data-directory variables once. + + Hoisted out of per-row loops on purpose: ``os.environ`` decodes the whole + environment on every iteration, which used to dominate database load time. + Returns None when the process is not linked to a primary MTDB.""" + try: + return {var: os.environ[var] for var in _PATH_COLUMNS.values()} + except KeyError: + return None + + +def rows2columns( + rows: Iterable[Mapping[str, Any]], + lineages: Mapping[str, Mapping[str, Any]], + add_paths: bool = True, +) -> Dict[str, list]: + """Build the column-oriented dict the `mtdb` class stores. + + Output is identical in shape to the flat reader's: one list per column, with + `taxonomy` holding a lineage dict per row that carries genus/species/strain + alongside the higher ranks.""" + df: Dict[str, list] = {c: [] for c in COLUMNS} + env = read_path_env() if add_paths else None + if add_paths and env is None: + logger.error("MycotoolsDB not in path, cannot delineate biofile paths") + + for row in rows: + row = dict(row) + ome = row.get("ome") or "" + genus = row.get("genus") or "" + # a fresh dict per row: callers mutate row taxonomy (df2db strips the + # genome-level ranks before writing) and must not corrupt the shared + # genus lineage + # reproduces `mtdb.read_tax`: stored ranks keep their order, then any + # standard rank the source omitted is appended empty + tax = dict(lineages.get(genus, {})) + for rank in _TAX_RANKS: + if rank not in tax: + tax[rank] = "" + tax["genus"] = genus + tax["species"] = (genus + " " + (row.get("species") or "")).strip() + tax["strain"] = row.get("strain") or "" + row["taxonomy"] = tax + _resolve_paths(row, ome, env) + for col in COLUMNS: + df[col].append(row.get(col, "")) + return df + + +def _abbreviate_paths(row: Mapping[str, Any], ome: str) -> Dict[str, str]: + """Collapse default file paths back to '' for storage. + + Mirrors what `mtdb.df2db` does when writing a flat file, so a database + stays valid when its root directory moves.""" + out = {} + env = read_path_env() + for col in _PATH_COLUMNS: + value = str(row.get(col) or "") + if env is not None: + default = env[_PATH_COLUMNS[col]] + ome + "." + col + if value == default: + value = "" + out[col] = value + return out + + +def split_taxonomy( + omes: Sequence[str], + genera: Sequence[str], + taxonomies: Sequence[Any], +) -> Dict[str, Dict[str, Any]]: + """Reduce per-row taxonomy to one lineage per genus. + + Lineages are a property of the genus -- `assimilate_tax` assigns the same + dict to every row of a genus -- so this is lossless for a well-formed + database. A genuine conflict means the source is corrupt, so it is reported + rather than silently resolved; the most complete lineage wins.""" + lineages: Dict[str, Dict[str, Any]] = {} + conflicts = set() + for ome, genus, tax in zip(omes, genera, taxonomies): + if not genus: + continue + stripped = _strip_genome_ranks(tax) + if not _rank_count(stripped): + lineages.setdefault(genus, stripped) + continue + prior = lineages.get(genus) + if prior is None or not _rank_count(prior): + lineages[genus] = stripped + elif prior != stripped: + conflicts.add(genus) + if _rank_count(stripped) > _rank_count(prior): + lineages[genus] = stripped + if conflicts: + logger.warning( + "%d genera carry conflicting lineages across rows (%s%s); kept the " + "most complete for each", + len(conflicts), + ", ".join(sorted(conflicts)[:5]), + "..." if len(conflicts) > 5 else "", + ) + return lineages + + +def _rank_count(lineage: Mapping[str, Any]) -> int: + return sum(1 for v in lineage.values() if v) + + +def _strip_genome_ranks(tax: Any) -> Dict[str, Any]: + """Drop genus/species/strain, which live in genome columns. + + Empty ranks are kept, and so is their order: a lineage stored here has to + reproduce the taxonomy field of the `.mtdb` it came from byte for byte, so + that flat -> SQLite -> flat is an identity.""" + if not tax: + return {} + if isinstance(tax, str): + try: + tax = json.loads(tax.replace("'", '"')) + except (json.JSONDecodeError, TypeError): + return {} + if not isinstance(tax, dict): + return {} + return { + rank: name + for rank, name in tax.items() + if rank not in {"genus", "species", "strain"} + } + + +# --------------------------------------------------------------------------- # +# writing +# --------------------------------------------------------------------------- # +def write_db(path: "str | Path", columns: Mapping[str, Sequence[Any]]) -> str: + """Write a whole MTDB to `path` atomically. + + The database is built in a sibling temporary file and renamed into place, so + readers -- including `primary_db()`, which picks a database off the + filesystem -- never observe a partially written database.""" + path = str(path) + tmp = path + ".tmp" + for stale in (tmp, tmp + "-journal", tmp + "-wal", tmp + "-shm"): + if Path(stale).exists(): + Path(stale).unlink() + + omes = list(columns.get("ome", [])) + lineages = split_taxonomy( + omes, list(columns.get("genus", [])), list(columns.get("taxonomy", [])) + ) + + conn = sqlite3.connect(tmp) + try: + _init_schema(conn) + genome_rows = [] + for i, ome in enumerate(omes): + if not ome: + continue + row = {c: columns[c][i] if c in columns else "" for c in _GENOME_COLUMNS} + row.update(_abbreviate_paths(row, ome)) + genome_rows.append( + tuple("" if row.get(c) is None else str(row.get(c, "")) for c in _GENOME_COLUMNS) + ) + placeholders = ", ".join("?" * len(_GENOME_COLUMNS)) + conn.executemany( + f"INSERT OR REPLACE INTO genome ({', '.join(_GENOME_COLUMNS)}) " + f"VALUES ({placeholders})", + genome_rows, + ) + conn.executemany( + "INSERT OR REPLACE INTO taxonomy (genus, lineage) VALUES (?, ?)", + [(genus, json.dumps(lin)) for genus, lin in lineages.items()], + ) + conn.commit() + conn.execute("VACUUM") + finally: + conn.close() + + Path(tmp).replace(path) + return path + + +def upsert_rows(path: "str | Path", rows: Iterable[Mapping[str, Any]]) -> None: + """Insert or update individual genome rows in an existing database.""" + conn = connect(path, read_only=False) + try: + _init_schema(conn) + for row in rows: + ome = row.get("ome") or "" + if not ome: + continue + stored = {c: "" if row.get(c) is None else str(row.get(c, "")) for c in _GENOME_COLUMNS} + stored.update(_abbreviate_paths(row, ome)) + conn.execute( + f"INSERT OR REPLACE INTO genome ({', '.join(_GENOME_COLUMNS)}) " + f"VALUES ({', '.join('?' * len(_GENOME_COLUMNS))})", + tuple(stored[c] for c in _GENOME_COLUMNS), + ) + lineage = _strip_genome_ranks(row.get("taxonomy")) + if lineage and row.get("genus"): + conn.execute( + "INSERT OR REPLACE INTO taxonomy (genus, lineage) VALUES (?, ?)", + (row["genus"], json.dumps(lineage)), + ) + conn.commit() + finally: + conn.close() + + +def delete_omes(path: "str | Path", omes: Iterable[str]) -> int: + """Remove genomes by ome; returns the number of rows deleted.""" + conn = connect(path, read_only=False) + try: + cur = conn.executemany("DELETE FROM genome WHERE ome = ?", [(o,) for o in omes]) + conn.commit() + return cur.rowcount + finally: + conn.close() + + +# --------------------------------------------------------------------------- # +# reading +# --------------------------------------------------------------------------- # +def _lineages(conn: sqlite3.Connection, genera: Optional[Iterable[str]] = None) -> Dict[str, dict]: + if genera is None: + rows = conn.execute("SELECT genus, lineage FROM taxonomy").fetchall() + else: + genera = list(dict.fromkeys(genera)) + rows = [] + for chunk in _chunks(genera): + rows.extend( + conn.execute( + f"SELECT genus, lineage FROM taxonomy WHERE genus IN " + f"({', '.join('?' * len(chunk))})", + chunk, + ).fetchall() + ) + out = {} + for row in rows: + try: + out[row["genus"]] = json.loads(row["lineage"]) + except json.JSONDecodeError: + logger.error("malformed lineage for genus %s", row["genus"]) + out[row["genus"]] = {} + return out + + +def _chunks(seq: Sequence[Any], size: int = 900): + """SQLite caps host parameters per statement (default 999).""" + for i in range(0, len(seq), size): + yield list(seq[i : i + size]) + + +def read_db(path: "str | Path", add_paths: bool = True) -> Dict[str, list]: + """Read an entire SQLite MTDB into the column dict the `mtdb` class holds.""" + conn = connect(path) + try: + rows = conn.execute( + f"SELECT {', '.join(_GENOME_COLUMNS)} FROM genome ORDER BY ome" + ).fetchall() + lineages = _lineages(conn) + finally: + conn.close() + return rows2columns(rows, lineages, add_paths=add_paths) + + +def select_omes( + path: "str | Path", omes: Iterable[str], add_paths: bool = True +) -> Dict[str, list]: + """Read only the named genomes -- an index seek per ome, not a full scan. + + This is the query behind `mtdb `, `mtdb accession`, and every other + tool that needs a handful of genomes out of the primary database.""" + omes = [o for o in dict.fromkeys(omes) if o] + if not omes: + return {c: [] for c in COLUMNS} + conn = connect(path) + try: + rows = [] + for chunk in _chunks(omes): + rows.extend( + conn.execute( + f"SELECT {', '.join(_GENOME_COLUMNS)} FROM genome WHERE ome IN " + f"({', '.join('?' * len(chunk))}) ORDER BY ome", + chunk, + ).fetchall() + ) + lineages = _lineages(conn, [r["genus"] for r in rows]) + finally: + conn.close() + return rows2columns(rows, lineages, add_paths=add_paths) + + +def select_ome_prefix( + path: "str | Path", prefix: str, add_paths: bool = True +) -> Dict[str, list]: + """Read the genome named `prefix`, or its MTDB-versioned successors. + + An ome may carry a version tag (`cryneo24` -> `cryneo24.1`), so a lookup for + the base ome has to find the versioned row.""" + conn = connect(path) + try: + rows = conn.execute( + f"SELECT {', '.join(_GENOME_COLUMNS)} FROM genome " + f"WHERE ome = ? OR ome LIKE ? ESCAPE '\\' ORDER BY ome", + (prefix, prefix.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_") + ".%"), + ).fetchall() + lineages = _lineages(conn, [r["genus"] for r in rows]) + finally: + conn.close() + return rows2columns(rows, lineages, add_paths=add_paths) + + +def select_column( + path: "str | Path", column: str, values: Iterable[str], add_paths: bool = True +) -> Dict[str, list]: + """Read the genomes whose `column` matches one of `values`.""" + if column not in set(_GENOME_COLUMNS): + raise KeyError(f"cannot select on {column}") + values = [v for v in dict.fromkeys(values) if v] + if not values: + return {c: [] for c in COLUMNS} + conn = connect(path) + try: + rows = [] + for chunk in _chunks(values): + rows.extend( + conn.execute( + f"SELECT {', '.join(_GENOME_COLUMNS)} FROM genome WHERE {column} IN " + f"({', '.join('?' * len(chunk))}) ORDER BY ome", + chunk, + ).fetchall() + ) + lineages = _lineages(conn, [r["genus"] for r in rows]) + finally: + conn.close() + return rows2columns(rows, lineages, add_paths=add_paths) + + +def omes(path: "str | Path") -> List[str]: + """List every ome without materializing the rest of the database.""" + conn = connect(path) + try: + return [r[0] for r in conn.execute("SELECT ome FROM genome ORDER BY ome")] + finally: + conn.close() + + +def count(path: "str | Path") -> int: + conn = connect(path) + try: + return conn.execute("SELECT COUNT(*) FROM genome").fetchone()[0] + finally: + conn.close() + + +def infer_rank(path: "str | Path", lineage: str) -> Optional[str]: + """Find the taxonomic rank a lineage name belongs to, using the taxonomy + table instead of scanning every genome. Returns None if unknown.""" + target = lineage.lower() + conn = connect(path) + try: + if conn.execute( + "SELECT 1 FROM genome WHERE LOWER(genus) = ? LIMIT 1", (target,) + ).fetchone(): + return "genus" + for row in conn.execute("SELECT lineage FROM taxonomy"): + try: + lin = json.loads(row[0]) + except json.JSONDecodeError: + continue + for rank, name in lin.items(): + if isinstance(name, str) and name.lower() == target: + return rank + if conn.execute( + "SELECT 1 FROM genome WHERE LOWER(genus || ' ' || species) = ? LIMIT 1", + (target,), + ).fetchone(): + return "species" + if conn.execute( + "SELECT 1 FROM genome WHERE LOWER(strain) = ? LIMIT 1", (target,) + ).fetchone(): + return "strain" + finally: + conn.close() + return None + + +def genera_for_lineages(path: "str | Path", lineages: Iterable[str]) -> Optional[List[str]]: + """Genera whose stored lineage matches any of `lineages` at any rank. + + Returns None when a name resolves to a genome-level rank (species/strain), + which the taxonomy table cannot answer -- the caller then filters in memory.""" + wanted = set(x.lower() for x in lineages if x) + if not wanted: + return [] + conn = connect(path) + try: + hits = set() + for row in conn.execute("SELECT genus, lineage FROM taxonomy"): + genus = row["genus"] + if genus and genus.lower() in wanted: + hits.add(genus) + continue + try: + lin = json.loads(row["lineage"]) + except json.JSONDecodeError: + continue + for name in lin.values(): + if isinstance(name, str) and name.lower() in wanted: + hits.add(genus) + break + # a genus with no taxonomy row can still be matched by name + for row in conn.execute("SELECT DISTINCT genus FROM genome"): + if row[0] and row[0].lower() in wanted: + hits.add(row[0]) + finally: + conn.close() + return sorted(hits) diff --git a/mycotools/lib/subcmd.py b/mycotools/lib/subcmd.py new file mode 100644 index 0000000..99628c2 --- /dev/null +++ b/mycotools/lib/subcmd.py @@ -0,0 +1,110 @@ +#! /usr/bin/env python3 +"""Shared factory for nested subcommand dispatchers. + +Mycotools groups related tools under a parent command (e.g. `mycotools +download`, `mtdb accession`). A group dispatcher forwards +` ...` to a submodule that parses its own `sys.argv`. + +Rather than reimplement the same forward-argv-and-delegate logic in every +group's `__init__.py`, each group builds a `Dispatcher` with its subcommand map +and help text, then exposes the dispatcher's `main`/`cli`. The one exception is +the `mtdb` root command, which carries extra base operations (link/unlink/list, +ome lookup) and stays bespoke.""" +import sys +import logging +import argparse +import importlib + +logger = logging.getLogger(__name__) + + +class Dispatcher: + """A nested subcommand dispatcher for a single command group. + + Parameters + ---------- + prog : str + Program name shown in help and prepended to the forwarded argv, e.g. + ``"mycotools download"``. + package : str + Importable package the submodules live in, e.g. + ``"mycotools.download"``. A resolved subcommand ``jgi`` is run by + importing ``mycotools.download.jgi`` and calling its ``cli()``. + subcommands : dict + Map of subcommand name/alias -> submodule name within ``package``. + description : str + Full ``argparse`` description (RawDescriptionHelpFormatter), typically + a header plus a formatted subcommand listing. + metavar : str + Positional metavar shown in usage (default ``"SUBCOMMAND"``). + arg_help : str + Help string for the positional (default ``"subcommand to run (see below)"``). + """ + + def __init__( + self, + prog, + package, + subcommands, + description, + metavar="SUBCOMMAND", + arg_help="subcommand to run (see below)", + ): + self.prog = prog + self.package = package + self.subcommands = subcommands + self.description = description + self.metavar = metavar + self.arg_help = arg_help + + def build_parser(self): + """Build the group parser. + + The subcommand positional is followed by a REMAINDER so the submodule's + arguments (including forwarded flags such as `-h`) pass through verbatim; + native argparse subparsers would intercept them.""" + parser = argparse.ArgumentParser( + prog=self.prog, + description=self.description, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + parser.add_argument( + "subcommand", nargs="?", metavar=self.metavar, help=self.arg_help + ) + parser.add_argument("rest", nargs=argparse.REMAINDER, help=argparse.SUPPRESS) + return parser + + def delegate(self, module_name, args): + """Run a submodule's CLI in-process and return its exit code. + + The submodule parses `sys.argv` itself, so its argv is swapped in for the + call and the `SystemExit` it raises (argparse errors, explicit exits) is + translated back to an exit code. Imports are deferred so a bare group + invocation stays light.""" + module = importlib.import_module(f"{self.package}.{module_name}") + saved_argv = sys.argv + sys.argv = [f"{self.prog} {module_name}"] + list(args) + try: + module.cli() + return 0 + except SystemExit as exc: + if exc.code is None: + return 0 + return exc.code if isinstance(exc.code, int) else 1 + finally: + sys.argv = saved_argv + + def main(self, argv=None): + if argv is None: + argv = sys.argv + parser = self.build_parser() + args = parser.parse_args(argv[1:]) + if args.subcommand in self.subcommands: + sys.exit(self.delegate(self.subcommands[args.subcommand], args.rest)) + if args.subcommand is not None: + logger.error(f"invalid subcommand: {args.subcommand}") + parser.print_help(sys.stderr) + sys.exit(1) + + def cli(self): + self.main(sys.argv) diff --git a/mycotools/manage_mtdb.py b/mycotools/manage_mtdb.py deleted file mode 100755 index 8cb36aa..0000000 --- a/mycotools/manage_mtdb.py +++ /dev/null @@ -1,125 +0,0 @@ -#! /usr/bin/env python3 - -import os -import sys -import argparse -from mycotools.lib.dbtools import loginCheck, primaryDB, mtdb, encrypt_pw -from mycotools.lib.kontools import format_path, read_json, collect_files - - -def rm_outdated(omes, yes=False): - """Remove outdated genomes after compiling them""" - - biofiles, to_del = [], [] - # compile the files - biofiles.extend( - [f"{os.environ['MYCOGFF3']}/{x}" for x in os.listdir(os.environ["MYCOGFF3"])] - ) - biofiles.extend( - [f"{os.environ['MYCOFAA']}/{x}" for x in os.listdir(os.environ["MYCOFAA"])] - ) - biofiles.extend( - [f"{os.environ['MYCOFNA']}/{x}" for x in os.listdir(os.environ["MYCOFNA"])] - ) - - # remove each biofile - for i in biofiles: - ome = None - ome_prep = os.path.basename(i) - if ome_prep.endswith(".gff3"): - ome = ome_prep[:-5] - elif ome_prep.endswith(".faa"): - ome = ome_prep[:-4] - elif ome_prep.endswith(".fna"): - ome = ome_prep[:-4] - else: # safer to preserve independent placements - continue - if ome not in omes: - to_del.append(i) - - if to_del: - if yes: - data = "y" - else: - data = input(f"\n{len(to_del)} omes to be deleted.\n" + "Continue [y/N]? ") - if data.lower() in {"yes", "y"}: - for i in to_del: - os.remove(i) - else: - raise KeyError("cache removal stopped") - - -def restrictions( - db, restr_list, mtdb_config=format_path("~/.mycotools/config.json"), yes=False -): - mtdb_config = read_json(mtdb_config) - restr_path = mtdb_config[mtdb_config["active"]]["MYCODB"] + "../log/failed.tsv" - - try: - with open(restr_path, "r") as raw: - restricted = [x.rstrip().split() for x in raw] - except FileNotFoundError: - restricted = [] - - accs = set(x[0] for x in restricted) - for r, s, reason in restr_list: - if s.lower() in {"ncbi", "jgi"} and r not in accs: - restricted.append([r, s.lower(), str(reason)]) - print(r, s, flush=True) - - in_db = [x[0] for x in restricted if x[0] in db] - while in_db: - if not yes: - check = input("Some restrictions are in the MTDB. Delete them? [y/N]: ") - if check.lower() in {"yes", "y"}: - break - else: - sys.exit(1) - - with open(restr_path, "w") as out: - out.write("\n".join(["\t".join(x) for x in restricted])) - - -def cli(): - parser = argparse.ArgumentParser( - description="Primary MycotoolsDB management utility" - ) - parser.add_argument( - "-c", "--clear_cache", action="store_true", help="Clear MycotoolsDB legacy data" - ) - parser.add_argument( - "-p", - "--password", - action="store_true", - help="Encrypt NCBI/JGI passwords to expedite access", - ) - parser.add_argument( - "-r", - "--restrict", - help="Restrict assembly accessions file, formatted: " - + "\t\t[REASON]", - ) - parser.add_argument("-y", "--yes", help="Answer yes", action="store_true") - args = parser.parse_args() - - db = mtdb(primaryDB()).set_index("assembly_acc") - - if args.password: - ncbi_email, ncbi_api, jgi_email, jgi_pwd = loginCheck() - encrypt_pw(ncbi_email, ncbi_api, jgi_email, jgi_pwd) - if args.restrict: - restrict_path = format_path(args.restrict) - with open(restrict_path, "r") as raw: - restricted = [x.rstrip().split("\t") for x in raw] - for v in restricted: - if len(v) < 3: - v = v + [None] - restrictions(db, restricted, yes=args.yes) - if args.clear_cache: - rm_outdated(mtdb(primaryDB())["ome"], args.yes) - - sys.exit(0) - - -if __name__ == "__main__": - cli() diff --git a/mycotools/mtdb.py b/mycotools/mtdb.py deleted file mode 100755 index c4f679e..0000000 --- a/mycotools/mtdb.py +++ /dev/null @@ -1,212 +0,0 @@ -#! /usr/bin/env python3 - -# NEED to add a log option -# list of DBs, ability to change names, quickly change between -# list update dates -# report storage information -# report taxonomy data -# NEED to remove standalone scripts from PATH and just reference mtdb (legacy) -# NEED to add option to export NCBI/JGI credentials -# NEED to pay attention to old ome versions - -import os -import re -import sys -import subprocess -from mycotools.lib.kontools import format_path, read_json, write_json, eprint -from mycotools.lib.dbtools import ( - primaryDB, - mtdb_connect, - mtdb_disconnect, - mtdb_initialize, - mtdb, - loginCheck, - parse_user_config, -) - - -def get_version(): - from importlib.metadata import version - - print(f'Mycotools version {version("mycotools")}') - - -def main(argv=sys.argv): - description = ( - "MycotoolsDB (MTDB) utility usage" - + "\n\nmtdb: print master MTDB path" - + "\n\nmtdb :" - + "\n[e]xtract\t\textract sub .mtdb file" - + "\n[u]pdate\t\tupdate/initialize primary MTDB" - + "\n[p]redb2mtdb\t\tadd local genomes to the primary MTDB" - + "\n[m]anage\t\tMTDB management utility" - + "\n\nmtdb [.gff3|.fna|.faa]: [PATH] print ome/ome code path" - + "\n\nmtdb :" - + "\n[-i DBPATH]\t[--interface]\tInitialize MTDB linkage" - + "\n[-u]\t\t[--unlink]\tUnlink from MTDB" - + "\n[-l]\t\t[--list]\tList historically linked primary MTDBs" - + "\n[-d]\t\t[--dependency]\tInstall/update dependencies" - + "\n[-v]\t\t[--version]\tPrint Mycotools version and exit" - ) - # + '\n[-f]\t\t[--fungi] \tConnect to fungal MTDB' \ - # + '\n[-p]\t\t[--prokaryote]\tConnect to prokaryote MTDB' \ - - config = parse_user_config() - if any([not x.startswith("-") for x in argv[1:]]): - script = [x for x in argv[1:] if not x.startswith("-")] - script = script[0] - ab2mt = { - "extract": "extract_mtdb", - "update": "update_mtdb", - "predb2mtdb": "predb2mtdb", - "e": "extract_mtdb", - "u": "update_mtdb", - "p": "predb2mtdb", - "m": "manage_mtdb", - "manage": "manage_mtdb", - } - if script in ab2mt: - exit_code = subprocess.call( - [ab2mt[script]] + [x for x in argv[1:] if x != script] - ) - sys.exit(exit_code) - - if len([x for x in argv if x.startswith("-")]) > 1: - eprint("\nERROR: one argument allowed.\n" + description, flush=True) - sys.exit(2) - - set_argv = set(argv) - if {"-h", "--help"}.intersection(set_argv): - print("\n" + description + "\n", flush=True) - sys.exit(0) - elif {"-v", "--version"}.intersection(set_argv): - get_version() - sys.exit(0) - - if {"-i", "--interface"}.intersection(set_argv): - if "-i" in set_argv: - coord = "-i" - else: - coord = "--interface" - try: - mycodb_loc_prep = argv[argv.index(coord) + 1] - if isinstance(mycodb_loc_prep, str): - if mycodb_loc_prep.startswith("-"): - raise IndexError - else: - mycodb_loc = format_path(mycodb_loc_prep) - if not os.path.isdir(mycodb_loc): - raise FileNotFoundError("invalid MTDB path") - elif not os.path.isdir(mycodb_loc + "mtdb"): - raise FileNotFoundError("invalid MTDB path") - except IndexError: - raise ValueError("--interface requires a path") - mtdb_initialize(mycodb_loc) - elif {"-d", "--dependencies"}.intersection(set_argv): - pip_deps = ["dna_features_viewer", "mycotools"] - dep_cmds = [ - ["conda", "install", "-y", "clipkit"], - ["python3", "-m", "pip", "install"] + pip_deps + ["--upgrade"], - ] - for dep_cmd in dep_cmds: - cmd = subprocess.call(dep_cmd) - if cmd: - print("\nUPDATE failed. Exit " + str(dep_cmd)) - sys.exit(cmd) - elif {"-l", "--list"}.intersection(set_argv): - if "log" in config: - for dbtype in config["log"]: - print(dbtype, flush=True) - for mtdb_loc, login_time in config["log"][dbtype].items(): - print(f"\t{mtdb_loc} {login_time}", flush=True) - print() - sys.exit(0) - - # elif {'-f', '--fungi'}.intersection(set_argv): - # if len(set_argv) > 2: - # eprint('\nERROR: -f does not accept additional arguments. Did you mean -i?') - # sys.exit(14) - # if '-f' in set_argv: - # coord = '-f' - # else: - # coord = '--fungi' - # if 'fungi' not in config: - # raise ValueError('fungal MTDB not connected') - # else: - # mtdb_connect(config, 'fungi') - # if not os.path.isfile(format_path('~/.mycotools/mtdb_key')): - # loginCheck() - # elif {'-p', '--prokaryote'}.intersection(set_argv): - # if len(set_argv) > 2: - # eprint('\nERROR: -p does not accept additional arguments. Did you mean -i?') - # sys.exit(15) - # - # if '-p' in set_argv: - # coord = '-p' - # else: - # coord = '--prokaryote' - # if 'prokaryote' not in config: - # raise ValueError('prokaryote MTDB not connected') - # else: - # mtdb_connect(config, 'prokaryote') - # if not os.path.isfile(format_path('~/.mycotools/mtdb_key')): - # loginCheck() - elif {"-u", "--unlink"}.intersection(set_argv): - if "-u" in set_argv: - coord = "-u" - else: - coord = "--unlink" - mtdb_disconnect() - sys.exit(0) - elif len(argv) > 1: - omes = argv[1].replace('"', "").replace("'", "").split() - for ome_prep in omes: - db = mtdb(primaryDB()).set_index() - if ome_prep in db: - print( - ome_prep - + "\t" - + "\t".join([str(db[ome_prep][x]) for x in db[ome_prep]]) - ) - sys.exit(0) - ome = re.sub(r"\.\w+[\w\d]$", "", ome_prep) - extension_srch = re.search(r"^\d+\.?\d*\.(.*$)", ome_prep[6:]) - if extension_srch is not None: - extension = extension_srch[1] - else: - extension = None - if ome in db: - try: - if extension: - print(db[ome][extension]) - else: - print({**{"ome": ome}, **db[ome]}) - except KeyError: - raise KeyError("Invalid extension " + extension) - else: - for ref_ome, row in db.items(): - if ref_ome.startswith(ome + "."): - if extension: - print(row[extension]) - else: - print({**{"ome": ome}, **row}) - break - else: - raise KeyError("Invalid ome " + ome) - sys.exit(0) - - path = primaryDB() - if path: - print(path, flush=True) - sys.exit(0) - else: - eprint("Link a MycotoolsDB via `mtdb -i `") - sys.exit(1) - - -def cli(): - main(sys.argv) - - -if __name__ == "__main__": - cli() diff --git a/mycotools/mtdb/__init__.py b/mycotools/mtdb/__init__.py new file mode 100755 index 0000000..3fc5d25 --- /dev/null +++ b/mycotools/mtdb/__init__.py @@ -0,0 +1,231 @@ +#! /usr/bin/env python3 + +# NEED to add a log option +# list of DBs, ability to change names, quickly change between +# list update dates +# report storage information +# report taxonomy data +# NEED to add option to export NCBI/JGI credentials +# NEED to pay attention to old ome versions + +import re +import sys +import logging +import argparse +import importlib +from pathlib import Path +from mycotools.lib.kontools import format_path, setup_logging +from mycotools.lib import mtdb_sql +from mycotools.lib.dbtools import ( + primary_db, + mtdb_disconnect, + mtdb_initialize, + mtdb, + parse_user_config, +) + +logger = logging.getLogger(__name__) + +# subcommand name/alias -> submodule within this package (mycotools.mtdb.). +# `accession`/`a` targets the acc2 subpackage, which further dispatches by format. +SUBCOMMANDS = { + "extract": "extract", + "e": "extract", + "update": "update", + "u": "update", + "predb": "predb", + "p": "predb", + "configure": "configure", + "config": "configure", + "c": "configure", + "manage": "manage", + "m": "manage", + "accession": "acc2", + "a": "acc2", + "files": "files", + "f": "files", +} + +DESCRIPTION = """MycotoolsDB (MTDB) utility + +Run without arguments to print the primary MTDB path. + +Subcommands (all following arguments are forwarded to the subcommand): + extract (e) extract a sub-.mtdb file + update (u) update / initialize the primary MTDB + predb (p) add local genomes to the primary MTDB + configure (c) change the primary MTDB configuration + manage (m) MTDB management utility + accession (a) retrieve data for accession(s) by format (fa/gff/gbk/locus) + files (f) symlink/copy selected files from the database + +Ome lookup: + mtdb [.gff3|.fna|.faa] print an ome's row, or a specific file path""" + + +def get_version(): + """Return the `-v`/`--version` output string.""" + from importlib.metadata import version + + return f'Mycotools version {version("mycotools")}' + + +def build_parser(): + """Build the base `mtdb` argument parser. + + Subcommands and ome-lookup share one positional (`target`) followed by a + REMAINDER: this lets subcommand arguments pass through verbatim to their + subcommand module (native argparse subparsers cannot, as they intercept + forwarded flags and reject arbitrary ome positionals).""" + parser = argparse.ArgumentParser( + prog="mtdb", + description=DESCRIPTION, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + parser.add_argument("-v", "--version", action="version", version=get_version()) + ops = parser.add_mutually_exclusive_group() + ops.add_argument( + "-i", "--interface", metavar="DBPATH", help="link/initialize the MTDB at path" + ) + ops.add_argument("-u", "--unlink", action="store_true", help="unlink from the MTDB") + ops.add_argument( + "-l", + "--list", + action="store_true", + help="list historically linked primary MTDBs", + ) + parser.add_argument( + "target", + nargs="?", + metavar="SUBCOMMAND|OME", + help="subcommand (see below) or ome code to look up", + ) + parser.add_argument("rest", nargs=argparse.REMAINDER, help=argparse.SUPPRESS) + return parser + + +def delegate(module_name, args): + """Run a subcommand's CLI in-process and return its exit code. + + The subcommand modules live in this package (`mycotools.mtdb.`) and + parse `sys.argv` themselves, so their argv is swapped in for the call and the + `SystemExit` they raise (argparse errors, explicit exits) is translated back + to an exit code. Imports are deferred so a bare `mtdb` invocation stays light + - `update` in particular pulls in the JGI/NCBI download stack.""" + module = importlib.import_module(f"mycotools.mtdb.{module_name}") + saved_argv = sys.argv + sys.argv = [f"mtdb {module_name}"] + list(args) + try: + module.cli() + return 0 + except SystemExit as exc: + if exc.code is None: + return 0 + return exc.code if isinstance(exc.code, int) else 1 + finally: + sys.argv = saved_argv + + +def link_mtdb(path): + """Link (interface) to the MTDB rooted at `path`.""" + mtdb_loc = format_path(path) + if not Path(mtdb_loc).is_dir() or not Path(mtdb_loc + "mtdb").is_dir(): + raise FileNotFoundError("invalid MTDB path") + mtdb_initialize(mtdb_loc) + + +def list_links(config): + """Print the historically linked primary MTDBs recorded in the user config.""" + for dbtype in config.get("log", {}): + print(dbtype, flush=True) + for mtdb_loc, login_time in config["log"][dbtype].items(): + print(f"\t{mtdb_loc} {login_time}", flush=True) + print() + + +def _resolve_ome(db_path, ome_prep): + """Return (ome, row) for a lookup token, or (None, None). + + Against a SQLite primary this is an index seek per token; against a `.mtdb` + flat file the whole database has to be parsed, so the load is deferred until + a flat file is actually what we have.""" + ome = re.sub(r"\.\w+[\w\d]$", "", ome_prep) + if mtdb_sql.is_sqlite(db_path): + for candidate in (ome_prep, ome): + rows = mtdb(mtdb_sql.select_ome_prefix(db_path, candidate)).set_index("ome") + if rows: + found = candidate if candidate in rows else sorted(rows)[0] + return found, rows[found] + return None, None + db = mtdb(db_path).set_index() + if ome_prep in db: + return ome_prep, db[ome_prep] + if ome in db: + return ome, db[ome] + for ref_ome, row in sorted(db.items()): + if ref_ome.startswith(ome + "."): + return ome, row + return None, None + + +def lookup_omes(omes): + """Print the database row, or a specific file path, for ome code(s).""" + db_path = primary_db() + for ome_prep in omes: + ext_srch = re.search(r"^\d+\.?\d*\.(.*$)", ome_prep[6:]) + extension = ext_srch[1] if ext_srch is not None else None + ome, row = _resolve_ome(db_path, ome_prep) + if row is None: + raise KeyError("Invalid ome " + ome_prep) + if extension: + if extension not in row: + raise KeyError("Invalid extension " + extension) + print(row[extension], flush=True) + else: + print( + ome + "\t" + "\t".join(str(v) for v in row.values()), + flush=True, + ) + + +def print_primary(): + """Print the primary MTDB path; return an exit code.""" + path = primary_db() + if path: + print(path, flush=True) + return 0 + logger.error("Link a MycotoolsDB via `mtdb -i `") + return 1 + + +def main(argv=sys.argv): + setup_logging() + config = parse_user_config() + args = build_parser().parse_args(argv[1:]) + + # 1. forward recognized subcommands to their standalone tool + if args.target in SUBCOMMANDS: + sys.exit(delegate(SUBCOMMANDS[args.target], args.rest)) + # 2. terminal base operations + if args.unlink: + mtdb_disconnect() + sys.exit(0) + if args.list: + list_links(config) + sys.exit(0) + # 3. ome-code lookup (skipped when linking, which reports the path afterward) + if args.target and not args.interface: + lookup_omes(args.target.replace('"', "").replace("'", "").split()) + sys.exit(0) + # 4. optionally (re)link, then always report the primary MTDB path + if args.interface: + link_mtdb(args.interface) + sys.exit(print_primary()) + + +def cli(): + main(sys.argv) + + +if __name__ == "__main__": + cli() diff --git a/mycotools/mtdb/__main__.py b/mycotools/mtdb/__main__.py new file mode 100644 index 0000000..6ad6bb9 --- /dev/null +++ b/mycotools/mtdb/__main__.py @@ -0,0 +1,6 @@ +#! /usr/bin/env python3 +"""Enable ``python -m mycotools.mtdb`` to run the MTDB dispatcher.""" +from mycotools.mtdb import cli + +if __name__ == "__main__": + cli() diff --git a/mycotools/mtdb/acc2/__init__.py b/mycotools/mtdb/acc2/__init__.py new file mode 100644 index 0000000..26f521c --- /dev/null +++ b/mycotools/mtdb/acc2/__init__.py @@ -0,0 +1,98 @@ +#! /usr/bin/env python3 +"""Accession-retrieval dispatcher for the `mtdb accession` subcommand. + +Routes `mtdb accession ...` (equivalently `mtdb a ...`) to a +per-format retrieval module in this package. Each format module parses its own +`sys.argv`, so its argv is swapped in and run in-process; deferred imports keep a +bare `mtdb accession` light.""" +import sys +import logging +import argparse +import importlib + +logger = logging.getLogger(__name__) + +# format name/alias -> submodule within this package (mycotools.mtdb.acc2.) +SUBCOMMANDS = { + "fa": "fa", + "fasta": "fa", + "gff": "gff", + "gff3": "gff", + "gbk": "gbk", + "genbank": "gbk", + "locus": "locus", + "l": "locus", +} + +DESCRIPTION = """Retrieve MTDB data for accession(s) in a given file format + +Formats (all following arguments are forwarded to the format's retriever): + fa (fasta) retrieve protein/nucleotide FASTA for accession(s) + gff (gff3) retrieve GFF3 entries for accession(s) + gbk (genbank) generate a GenBank file for accession(s)/ome(s) + locus (l) retrieve the locus surrounding accession(s) + +Examples: + mtdb accession fa -a _ + mtdb accession gff -h""" + + +def build_parser(): + """Build the `mtdb accession` argument parser. + + The format positional is followed by a REMAINDER so the format module's + arguments pass through verbatim (native argparse subparsers would intercept + forwarded flags such as `-h`).""" + parser = argparse.ArgumentParser( + prog="mtdb accession", + description=DESCRIPTION, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + parser.add_argument( + "format", + nargs="?", + metavar="FORMAT", + help="retrieval format (see below)", + ) + parser.add_argument("rest", nargs=argparse.REMAINDER, help=argparse.SUPPRESS) + return parser + + +def delegate(module_name, args): + """Run a format retriever's CLI in-process and return its exit code. + + The format modules live in this package (`mycotools.mtdb.acc2.`) and + parse `sys.argv` themselves, so their argv is swapped in for the call and the + `SystemExit` they raise (argparse errors, explicit exits) is translated back + to an exit code.""" + module = importlib.import_module(f"mycotools.mtdb.acc2.{module_name}") + saved_argv = sys.argv + sys.argv = [f"mtdb accession {module_name}"] + list(args) + try: + module.cli() + return 0 + except SystemExit as exc: + if exc.code is None: + return 0 + return exc.code if isinstance(exc.code, int) else 1 + finally: + sys.argv = saved_argv + + +def main(argv=sys.argv): + parser = build_parser() + args = parser.parse_args(argv[1:]) + if args.format in SUBCOMMANDS: + sys.exit(delegate(SUBCOMMANDS[args.format], args.rest)) + if args.format is not None: + logger.error(f"invalid format: {args.format}") + parser.print_help(sys.stderr) + sys.exit(1) + + +def cli(): + main(sys.argv) + + +if __name__ == "__main__": + cli() diff --git a/mycotools/mtdb/acc2/__main__.py b/mycotools/mtdb/acc2/__main__.py new file mode 100644 index 0000000..80aef6d --- /dev/null +++ b/mycotools/mtdb/acc2/__main__.py @@ -0,0 +1,6 @@ +#! /usr/bin/env python3 +"""Enable ``python -m mycotools.mtdb.acc2`` to run the accession dispatcher.""" +from mycotools.mtdb.acc2 import cli + +if __name__ == "__main__": + cli() diff --git a/mycotools/acc2fa.py b/mycotools/mtdb/acc2/fa.py similarity index 88% rename from mycotools/acc2fa.py rename to mycotools/mtdb/acc2/fa.py index 8167a4d..f5bbc09 100755 --- a/mycotools/acc2fa.py +++ b/mycotools/mtdb/acc2/fa.py @@ -1,13 +1,15 @@ #! /usr/bin/env python3 -import os +import logging import re import sys import argparse from collections import defaultdict from mycotools.lib.biotools import fa2dict, dict2fa, reverse_complement -from mycotools.lib.dbtools import mtdb, primaryDB -from mycotools.lib.kontools import format_path, eprint, stdin2str +from mycotools.lib.dbtools import primary_db, load_omes, omes_from_accessions +from mycotools.lib.kontools import format_path, stdin2str, setup_logging + +logger = logging.getLogger(__name__) def extract_mtdb_accs_exp(fa_dict, accs): @@ -34,7 +36,7 @@ def extract_mtdb_accs(fa_dict, accs, spacer=""): try: out_fa[acc] = fa_dict[acc] except KeyError: - eprint(f"WARNING: {acc} has no CDS", flush=True) + logger.warning(f"{acc} has no CDS") continue acc_name = acc[: acc.find("[")] if start < end: @@ -44,9 +46,7 @@ def extract_mtdb_accs(fa_dict, accs, spacer=""): "description": fa_dict[acc_name]["description"], } except KeyError: - eprint( - spacer + "WARNING: invalid accession " + acc_name, flush=True - ) + logger.warning(spacer + "invalid accession " + acc_name) else: try: out_fa[acc] = { @@ -56,19 +56,17 @@ def extract_mtdb_accs(fa_dict, accs, spacer=""): "description": fa_dict[acc_name]["description"], } except KeyError: - eprint( - spacer + "WARNING: invalid accession " + acc_name, flush=True - ) + logger.warning(spacer + "invalid accession " + acc_name) else: # no coordinates try: out_fa[acc] = fa_dict[acc] # extract the whole accession except KeyError: - eprint(spacer + "WARNING: invalid accession " + acc_name, flush=True) + logger.warning(spacer + "invalid accession " + acc_name) return out_fa -def extractHeaders(fasta_file, accessions, ome=None): +def extract_headers(fasta_file, accessions, ome=None): """searches headers for "[]", which indicate coordinate-based extraction. otherwise, just retrieves the accession from the fasta dictionary""" @@ -138,7 +136,7 @@ def dbmain(db, accs, error=True, spacer="\t\t\t", coord_check=True): if error: raise KeyError("invalid ome: " + ome) else: - eprint(spacer + ome + " not in database", flush=True) + logger.info(spacer + ome + " not in database") fa_dict = {**fa_dict, **extract_mtdb_accs(ome_fasta, ome_accs)} else: for ome, ome_accs in ome_data.items(): @@ -148,7 +146,7 @@ def dbmain(db, accs, error=True, spacer="\t\t\t", coord_check=True): if error: raise KeyError else: - eprint(spacer + ome + " not in database", flush=True) + logger.info(spacer + ome + " not in database") fa_dict = {**fa_dict, **extract_mtdb_accs_exp(ome_fasta, ome_accs)} return fa_dict @@ -158,7 +156,7 @@ def famain(accs, fa, ome=None): """takes in accessions, fasta, and retrieves accessions""" fa_dict = {} - fa_dict = {**fa_dict, **extractHeaders(fa, accs, ome)} + fa_dict = {**fa_dict, **extract_headers(fa, accs, ome)} return fa_dict @@ -185,8 +183,9 @@ def cli(): "-s", "--start", help="Start index column (1 indexed)", type=int ) parser.add_argument("-e", "--end", help="End index column (1 indexed)", type=int) - parser.add_argument("-d", "--mtdb", default=primaryDB()) + parser.add_argument("-d", "--mtdb", default=primary_db()) args = parser.parse_args() + setup_logging(verbose=getattr(args, "verbose", False)) if args.input: # input file input_file = format_path(args.input) @@ -227,7 +226,8 @@ def cli(): db_path = format_path(args.mtdb) if not args.fasta: # MTDB run - db = mtdb(db_path) + # only the genomes behind the requested accessions are needed + db = load_omes(db_path, omes_from_accessions(accs)) fa_dict = dbmain(db, accs) fasta_str = dict2fa(fa_dict) else: # non MTDB diff --git a/mycotools/acc2gbk.py b/mycotools/mtdb/acc2/gbk.py similarity index 93% rename from mycotools/acc2gbk.py rename to mycotools/mtdb/acc2/gbk.py index c228ee1..e1f51e6 100755 --- a/mycotools/acc2gbk.py +++ b/mycotools/mtdb/acc2/gbk.py @@ -1,19 +1,20 @@ #! /usr/bin/env python3 -import os +import logging import re import sys import argparse -import multiprocessing as mp from itertools import chain from collections import defaultdict -from mycotools.lib.kontools import eprint, format_path, stdin2str -from mycotools.lib.dbtools import mtdb, primaryDB -from mycotools.lib.biotools import fa2dict, gff2list, gff3Comps -from mycotools.acc2gff import db_main as acc2gff +from mycotools.lib.kontools import format_path, stdin2str, setup_logging +from mycotools.lib.dbtools import mtdb, primary_db, load_omes, omes_from_accessions +from mycotools.lib.biotools import fa2dict, gff2list, gff3_comps +from mycotools.mtdb.acc2.gff import db_main as acc2gff +logger = logging.getLogger(__name__) -def col_CDS( + +def col_cds( gff_list, types={"gene", "CDS", "exon", "mRNA", "tRNA", "rRNA", "RNA", "pseudogene"} ): """Collect all CDS entries from a `gff` and store them into cds_dict. @@ -26,9 +27,9 @@ def col_CDS( if entry["type"] in types: contig = entry["seqid"] try: - alias = re.search(gff3Comps()["Alias"], entry["attributes"])[1] + alias = re.search(gff3_comps()["Alias"], entry["attributes"])[1] except TypeError: - eprint("\n\tERROR: could not extract Alias ID from " + gff, flush=True) + logger.error("could not extract Alias ID from " + gff) continue aliases = alias.split("|") # to address alternate splicing in gene # aliases @@ -89,7 +90,6 @@ def contig2gbk( + "+" ) name = ome + "_" + contig - relative_end = seq_coords[-1][1] - seq_coords[0][0] gbk = ( "LOCUS " + name @@ -146,14 +146,14 @@ def contig2gbk( # for each gene, let it be the parent entry for entry in entries["gene"]: - alias = re.search(gff3Comps()["Alias"], entry["attributes"])[1] + alias = re.search(gff3_comps()["Alias"], entry["attributes"])[1] # eprint(alias) if "|" in alias: # alternately spliced gene if alias in used_aliases: continue else: used_aliases.add(alias) - id_ = re.search(gff3Comps()["id"], entry["attributes"])[1] + id_ = re.search(gff3_comps()["id"], entry["attributes"])[1] products = {} # try to acquire the product name @@ -204,8 +204,8 @@ def contig2gbk( products[prod_id] = product except TypeError: # no product pass - alias = re.search(gff3Comps()["Alias"], entry["attributes"])[1] - id_ = re.search(gff3Comps()["par"], entry["attributes"])[1] + alias = re.search(gff3_comps()["Alias"], entry["attributes"])[1] + id_ = re.search(gff3_comps()["par"], entry["attributes"])[1] # append to the final gene coordinates if final_coords: @@ -279,8 +279,8 @@ def contig2gbk( final_coords += "\n " + cds_coords_list[-1] if len(cds_coords_list) > 1: final_coords = "join(" + final_coords + ")" - alias = re.search(gff3Comps()["Alias"], entry["attributes"])[1] - id_ = re.search(gff3Comps()["id"], entry["attributes"])[1] + alias = re.search(gff3_comps()["Alias"], entry["attributes"])[1] + id_ = re.search(gff3_comps()["id"], entry["attributes"])[1] products = {} for prod_id, product_search in product_searches.items(): @@ -377,8 +377,6 @@ def gen_gbk( if count0: if v1[0] - v0[1] > break_contigs: breaks.add(k0) - k0 = k1 - v0 = v1 count0 += 1 count1 = 0 @@ -451,7 +449,7 @@ def ome_main( faa, fna = fa2dict(row["faa"]), fa2dict(row["fna"]) for key, gffs in gff_lists.items(): for gff_list in gffs: - cds_dict = col_CDS( + cds_dict = col_cds( gff_list, types={ "gene", @@ -480,7 +478,7 @@ def main( break_contigs=False, ): """Generate a genbank for each inputed gff, its associated fna, and ome""" - cds_dict = col_CDS( + cds_dict = col_cds( gff_list, types={"gene", "CDS", "exon", "mRNA", "tRNA", "rRNA", "RNA", "pseudogene"}, ) @@ -532,9 +530,10 @@ def cli(): type=int, help="Base pairs to split long breaks between genes", ) - parser.add_argument("-d", "--mtdb", help="DEFAULT: master", default=primaryDB()) + parser.add_argument("-d", "--mtdb", help="DEFAULT: master", default=primary_db()) parser.add_argument("-c", "--cpu", type=int, default=1) args = parser.parse_args() + setup_logging(verbose=getattr(args, "verbose", False)) # gather accessions from various input styles and split into a list if args.input: @@ -566,7 +565,7 @@ def cli(): # only full genomes for args.full, no accessions if args.full: if any("_" in x for x in accs): # _ is forbidden from ome codes - eprint("\nERROR: -f requires omes, not accessions", flush=True) + logger.error("-f requires omes, not accessions") sys.exit(3) # create a default regular expression for the product name @@ -581,8 +580,14 @@ def cli(): for char in args.regex: regex += char - # import database and set index - db = mtdb(format_path(args.mtdb)) + # import database and set index; an inputted gff is not tied to the + # accession list, so only an MTDB run can narrow the read to the genomes + # actually referenced (omes directly for --full, else the accessions' omes) + db_path = format_path(args.mtdb) + if args.gff: + db = mtdb(db_path) + else: + db = load_omes(db_path, set(accs) if args.full else omes_from_accessions(accs)) db = db.set_index() # various output formats diff --git a/mycotools/acc2gff.py b/mycotools/mtdb/acc2/gff.py similarity index 87% rename from mycotools/acc2gff.py rename to mycotools/mtdb/acc2/gff.py index fefecfa..6e125dc 100755 --- a/mycotools/acc2gff.py +++ b/mycotools/mtdb/acc2/gff.py @@ -1,13 +1,16 @@ #! /usr/bin/env python3 -import os +import logging import re import sys import argparse import multiprocessing as mp from mycotools.lib.biotools import gff2list, list2gff -from mycotools.lib.dbtools import mtdb, primaryDB -from mycotools.lib.kontools import format_path, stdin2str +from mycotools.lib.dbtools import primary_db, load_omes, omes_from_accessions +from mycotools.lib.kontools import format_path, stdin2str, setup_logging +from pathlib import Path + +logger = logging.getLogger(__name__) def grab_gff_acc(gff_list, acc, term="Alias="): @@ -103,10 +106,11 @@ def cli(): "-o", "--ome", action="store_true", help="Output files by ome code" ) parser.add_argument( - "-d", "--mtdb", default=primaryDB(), help="mycodb DEFAULT: master" + "-d", "--mtdb", default=primary_db(), help="mycodb DEFAULT: master" ) parser.add_argument("--cpu", type=int, default=mp.cpu_count()) args = parser.parse_args() + setup_logging(verbose=getattr(args, "verbose", False)) # if there is an input file, extract the accessions from that if args.input: @@ -131,19 +135,15 @@ def cli(): else: accs = [args.accession] - if args.cpu < mp.cpu_count(): - cpu = args.cpu - else: - cpu = mp.cpu_count() # if no gff is provided, then acquire it from the primary database db_path = format_path(args.mtdb) if not args.gff: - db = mtdb(format_path(args.mtdb)) + # only the genomes behind the requested accessions are needed + db = load_omes(db_path, omes_from_accessions(accs)) gff_lists = db_main(db, accs, cpus=args.cpu) # otherwise just use what is available else: - gff_path = format_path(args.gff) gff_lists = gff_main(gff_data, accs) # if there is an inputted accession, then print the output to stdout @@ -151,13 +151,13 @@ def cli(): print(list2gff(gff_lists[list(gff_lists.keys())[0]]).rstrip(), flush=True) # if it is specified output, open a folder for it elif args.ome: - output = mkOutput(os.getcwd() + "/", "acc2gff") + output = mk_output(str(Path.cwd()) + "/", "acc2gff") for ome in gff_strs: if gff_lists[ome]: with open(output + ome + ".accs.gff3", "w") as out: out.write(list2gff(gff_lists[ome])) else: - eprint("ERROR: " + ome + " failed, no accessions retrieved", flush=True) + logger.error("" + ome + " failed, no accessions retrieved") # print to stdout each gff_list else: out_str = "" @@ -165,7 +165,7 @@ def cli(): if gff_lists[ome]: out_str += list2gff(gff_lists[ome]) + "\n" else: - eprint("ERROR: " + ome + " does not have accession", flush=True) + logger.error("" + ome + " does not have accession") print(out_str) sys.exit(0) diff --git a/mycotools/acc2locus.py b/mycotools/mtdb/acc2/locus.py similarity index 89% rename from mycotools/acc2locus.py rename to mycotools/mtdb/acc2/locus.py index 16b380f..1b56cf8 100755 --- a/mycotools/acc2locus.py +++ b/mycotools/mtdb/acc2/locus.py @@ -1,16 +1,18 @@ #! /usr/bin/env python3 -import os +import logging import re import sys import argparse import multiprocessing as mp from itertools import chain from collections import defaultdict -from mycotools.lib.kontools import eprint, format_path, file2list, stdin2str -from mycotools.lib.dbtools import primaryDB, mtdb -from mycotools.lib.biotools import gff2list, fa2dict, dict2fa, list2gff, gff3Comps -from mycotools.acc2gff import grab_gff_acc +from mycotools.lib.kontools import format_path, file2list, stdin2str, setup_logging +from mycotools.lib.dbtools import primary_db, load_omes, omes_from_accessions +from mycotools.lib.biotools import gff2list, fa2dict, dict2fa, list2gff, gff3_comps +from mycotools.mtdb.acc2.gff import grab_gff_acc + +logger = logging.getLogger(__name__) def prep_gff_output(hit_list, gff_path, cpu=1): @@ -46,7 +48,7 @@ def compile_alias_coords(gff_list, accs_list=[]): dictionary that is accessed through the sequence ID, followed by the alias of each sequence""" accs_set = set(accs_list) - alias_comp = re.compile(gff3Comps()["Alias"]) + alias_comp = re.compile(gff3_comps()["Alias"]) # gather the coordinates for each RNA and genes without RNAs coord_dict = defaultdict(lambda: defaultdict(list)) @@ -93,7 +95,7 @@ def compile_alias_coords(gff_list, accs_list=[]): return coord_dict, acc2seqid -def prep_outputXgene(coords_dict, acc, plusminus): +def prep_output_xgene(coords_dict, acc, plusminus): """Prep the output for each gene accession using the coordinates dictionary as a sorting mechanism, and return the list of accession names""" alias_list = list(coords_dict.keys()) @@ -108,11 +110,10 @@ def prep_outputXgene(coords_dict, acc, plusminus): return out_index -def prep_outputXbase(coords_dict, acc, plusminus): +def prep_output_xbase(coords_dict, acc, plusminus): """Prep the output based on the coordinates of a list of accessions if they are within the range of the bases alotted, provided by plusminus""" alias_list = list(coords_dict.keys()) - index = alias_list.index(acc) start, end = coords_dict[acc][0], coords_dict[acc][1] low_bound, high_bound = start - plusminus, end + plusminus # acquire the indices of the accessions that fit the boundary @@ -164,10 +165,10 @@ def main( out_indices[accs[0]] = grab_between(coords_dict[seqid], accs) elif nt: # if looking for accessions that are +/- a number of nucleotides for acc, seqid in acc2seqid.items(): - out_indices[acc] = prep_outputXbase(coords_dict[seqid], acc, plusminus) + out_indices[acc] = prep_output_xbase(coords_dict[seqid], acc, plusminus) else: # if looking for accessions that are +/- a number of accessions for acc, seqid in acc2seqid.items(): - out_indices[acc] = prep_outputXgene(coords_dict[seqid], acc, plusminus) + out_indices[acc] = prep_output_xgene(coords_dict[seqid], acc, plusminus) out_indices = {k: v for k, v in out_indices.items() if v} if geneGff: # if a gff of the RNA entries is desired @@ -177,7 +178,7 @@ def main( for entry in gff_list: if "RNA" in entry["type"]: try: - gene = re.search(gff3Comps()["Alias"], entry["attributes"])[1] + gene = re.search(gff3_comps()["Alias"], entry["attributes"])[1] for acc, genes in gene_sets.items(): if gene in genes: geneGffs_prep[acc][gene] = entry @@ -186,7 +187,7 @@ def main( pass elif "gene" in entry["type"]: try: - gene = re.search(gff3Comps()["Alias"], entry["attributes"])[1] + gene = re.search(gff3_comps()["Alias"], entry["attributes"])[1] for acc, genes in gene_sets.items(): if gene in genes: alt_geneGffs_prep[acc][gene] = entry @@ -258,15 +259,12 @@ def cli(): parser.add_argument("-f", "--faa", help="Input protein fasta file") parser.add_argument("-s", "--sep", help="Separator for input file.", default="\n") parser.add_argument( - "-d", "--mtdb", default=primaryDB(), help="MTDB; DEFAULT: primary" + "-d", "--mtdb", default=primary_db(), help="MTDB; DEFAULT: primary" ) parser.add_argument("--cpu", type=int, default=1) args = parser.parse_args() + setup_logging(verbose=getattr(args, "verbose", False)) - if args.cpu < mp.cpu_count(): - cpu = args.cpu - else: - cpu = mp.cpu_count() args.sep = args.sep.replace("'", "").replace('"', "") if args.input: @@ -284,26 +282,28 @@ def cli(): else: accs = [args.acc] else: - eprint("\nERROR: requires input or acc", flush=True) + logger.error("requires input or acc") sys.exit(1) if args.between: if len(accs) > 2: - eprint("\nERROR: -b needs 2 accessions", flush=True) + logger.error("-b needs 2 accessions") sys.exit(2) if args.nucleotide: - eprint("\nERROR: -b and -n are incompatible", flush=True) + logger.error("-b and -n are incompatible") sys.exit(3) db = None out_indices = {} + # only the genomes behind the requested accessions are needed + needed_omes = omes_from_accessions(accs) if args.gff: gff = gff2list(format_path(args.gff)) out_indices = main( gff, accs, args.plusminus, between=args.between, nt=args.nucleotide ) else: - db = mtdb(format_path(args.mtdb)).set_index("ome") + db = load_omes(format_path(args.mtdb), needed_omes).set_index("ome") out_indices = mycotools_main( db, accs, @@ -315,7 +315,7 @@ def cli(): if args.output: if not db: - db = mtdb(format_path(args.mtdb)).set_index("ome") + db = load_omes(format_path(args.mtdb), needed_omes).set_index("ome") for acc in out_indices: if args.gff: gff = format_path(args.gff) diff --git a/mycotools/mtdb/configure.py b/mycotools/mtdb/configure.py new file mode 100644 index 0000000..2e0543e --- /dev/null +++ b/mycotools/mtdb/configure.py @@ -0,0 +1,263 @@ +#! /usr/bin/env python3 +"""Modify a linked primary MycotoolsDB's persistent configuration. + +The options that shape what future `mtdb update` runs acquire -- inclusion of +use-restricted data, MycoCosm (JGI) participation, and lineage constraints -- +are established when the database is initialized (`mtdb update -i`). Because +those settings define the database, `mtdb update` only accepts them at +initialization; this utility is the supported way to change them afterward. + +It edits `config/mtdb.json` in place. With no options it prints the current +configuration. Every change takes effect on the next `mtdb update`.""" + +import os +import sys +import logging +import argparse +from collections import defaultdict +from pathlib import Path +from mycotools.lib.kontools import ( + format_path, + read_json, + write_json, + split_input, + setup_logging, +) + +logger = logging.getLogger(__name__) + +# ranks accepted for lineage constraints; mirrors `mtdb update` +PERMITTED_RANKS = {"phylum", "subphylum", "class", "order", "family", "genus"} + +# config keys written by update.gen_config, with a label for display. Ordered +# most- to least- frequently adjusted. +CONFIG_LABELS = ( + ("branch", "Kingdom/branch"), + ("jgi", "MycoCosm (JGI)"), + ("nonpublished", "Use-restricted data"), + ("lineage_constraints", "Lineage constraints"), + ("rogue", "Standalone (rogue)"), + ("repository", "Reference repository"), + ("forbidden", "Forbidden-ome ledger"), +) + +# values `nonpublished` may take when enabled (config stores "yes"; the historic +# forms are accepted defensively) +_TRUE = {"yes", "y", "true"} + + +def load_config(): + """Return (config_path, config) for the linked primary MTDB, or exit. + + The configuration lives beside the linked database, so a missing MYCODB is + "nothing is linked" and a missing file is a corrupt install -- the two + exits mirror the codes `mtdb update` raises for the same conditions.""" + if "MYCODB" not in os.environ: + logger.error("MTDB not linked. Link via `mtdb -i `") + sys.exit(50) + path = format_path("$MYCODB/../config/mtdb.json") + if not Path(path).is_file(): + logger.error("corrupted MycotoolsDB - no configuration found") + sys.exit(21) + return path, read_json(path) + + +def _is_true(value): + """Whether a stored `nonpublished`/`jgi` value counts as enabled.""" + if isinstance(value, str): + return value.lower() in _TRUE + return bool(value) + + +def fmt_value(key, value): + """Render a config value for display.""" + if key == "lineage_constraints": + if not value: + return "none" + return "; ".join( + f"{rank}: {', '.join(sorted(lineages))}" + for rank, lineages in sorted(value.items()) + ) + if key in {"nonpublished", "jgi", "rogue"}: + return "yes" if _is_true(value) else "no" + if value in (None, ""): + return "none" + return str(value) + + +def show_config(config): + """Print the current configuration.""" + print("MycotoolsDB configuration:", flush=True) + for key, label in CONFIG_LABELS: + if key in config: + print(f" {label}: {fmt_value(key, config[key])}", flush=True) + + +def parse_lineage_constraints(lineage, rank): + """Turn positional --lineage/--rank strings into a rank->[lineages] dict. + + Mirrors the parsing in `mtdb update`'s control_flow so a constraint set + reads identically whether it is established at init or reconfigured here. + Exits on a length mismatch or an unrecognized rank.""" + lineage_constraints = split_input(lineage) + rank_constraints = split_input(rank) + if len(lineage_constraints) != len(rank_constraints): + logger.error("--lineage must be same length as --rank") + sys.exit(18) + rank2lineages = defaultdict(set) + for i, lin in enumerate(lineage_constraints): + rank_c = rank_constraints[i].lower() + if rank_c not in PERMITTED_RANKS: + logger.error(f"accepted ranks: {sorted(PERMITTED_RANKS)}") + sys.exit(22) + rank2lineages[rank_c].add(lin.lower()) + return {k: sorted(v) for k, v in sorted(rank2lineages.items())} + + +def _set_nonpublished(config, args, changes): + """Toggle use-restricted data inclusion in `config`.""" + if args.nonpublished and args.published: + logger.error("--nonpublished and --published are mutually exclusive") + sys.exit(1) + if args.nonpublished: + if _is_true(config.get("nonpublished")): + logger.info("use-restricted data already enabled") + return + # the T&C acknowledgement lives in update; import it lazily so display + # and the other toggles do not pull in the download stack + from mycotools.mtdb.update import validate_t_and_c + + config["nonpublished"] = validate_t_and_c(config, discrepancy=True) + changes.append("use-restricted data enabled") + elif args.published: + if not _is_true(config.get("nonpublished")): + logger.info("use-restricted data already disabled") + return + config["nonpublished"] = False + changes.append("use-restricted data disabled") + + +def _set_jgi(config, args, changes): + """Toggle MycoCosm (JGI) participation in `config`.""" + if args.ncbi_only and args.jgi: + logger.error("--ncbi_only and --jgi are mutually exclusive") + sys.exit(1) + if args.ncbi_only: + if not _is_true(config.get("jgi")): + logger.info("MycoCosm already disabled") + return + config["jgi"] = False + changes.append("MycoCosm disabled (NCBI only)") + elif args.jgi: + if config.get("branch", "fungi") != "fungi": + logger.warning( + f'branch "{config.get("branch")}" does not use MycoCosm; ' + "enabling it has no effect" + ) + if _is_true(config.get("jgi")): + logger.info("MycoCosm already enabled") + return + config["jgi"] = True + changes.append("MycoCosm enabled") + + +def _set_lineage(config, args, changes): + """Set or clear the lineage constraints in `config`.""" + if args.clear_lineage and (args.lineage or args.rank): + logger.error("--clear_lineage cannot be combined with --lineage/--rank") + sys.exit(1) + if args.clear_lineage: + if config.get("lineage_constraints"): + config["lineage_constraints"] = {} + changes.append("lineage constraints cleared") + else: + logger.info("no lineage constraints to clear") + elif args.lineage or args.rank: + if not (args.lineage and args.rank): + logger.error("--lineage requires --rank") + sys.exit(16) + new = parse_lineage_constraints(args.lineage, args.rank) + if new != config.get("lineage_constraints"): + config["lineage_constraints"] = new + changes.append("lineage constraints updated") + else: + logger.info("lineage constraints unchanged") + + +def apply_changes(config, args): + """Mutate `config` in place per `args`; return a list of change summaries.""" + changes = [] + _set_nonpublished(config, args, changes) + _set_jgi(config, args, changes) + _set_lineage(config, args, changes) + return changes + + +def cli(): + parser = argparse.ArgumentParser( + description="Modify the linked primary MycotoolsDB configuration " + "(config/mtdb.json). With no options, print the current configuration. " + "Changes take effect on the next `mtdb update`." + ) + restr = parser.add_argument_group("Use-restricted data") + restr.add_argument( + "--nonpublished", + action="store_true", + help="[FUNGI]: Include MycoCosm use-restricted data", + ) + restr.add_argument( + "--published", + action="store_true", + help="Exclude use-restricted data", + ) + + jgi = parser.add_argument_group("MycoCosm") + jgi.add_argument( + "--ncbi_only", action="store_true", help="[FUNGI]: Forego MycoCosm (NCBI only)" + ) + jgi.add_argument("--jgi", action="store_true", help="[FUNGI]: Include MycoCosm") + + lin = parser.add_argument_group("Lineage constraints") + lin.add_argument("-l", "--lineage", help="Lineage(s) to constrain updates to") + lin.add_argument( + "-rk", "--rank", help="Rank(s) that positionally correspond to -l" + ) + lin.add_argument( + "--clear_lineage", action="store_true", help="Remove all lineage constraints" + ) + args = parser.parse_args() + setup_logging(verbose=getattr(args, "verbose", False)) + + path, config = load_config() + + requested = ( + args.nonpublished + or args.published + or args.ncbi_only + or args.jgi + or args.lineage + or args.rank + or args.clear_lineage + ) + if not requested: + show_config(config) + sys.exit(0) + + changes = apply_changes(config, args) + if changes: + write_json(config, path) + for change in changes: + logger.info(change) + logger.info("Run `mtdb update` to apply") + else: + logger.info("No configuration changes") + show_config(config) + sys.exit(0) + + +def main(): + cli() + + +if __name__ == "__main__": + cli() diff --git a/mycotools/extract_mtdb.py b/mycotools/mtdb/extract.py similarity index 51% rename from mycotools/extract_mtdb.py rename to mycotools/mtdb/extract.py index d8157bd..be7766f 100755 --- a/mycotools/extract_mtdb.py +++ b/mycotools/mtdb/extract.py @@ -4,120 +4,44 @@ # NEED stdin acceptance for most of these arguments import os -import re import sys -import copy -import random +import logging import argparse -from collections import defaultdict from mycotools.lib.kontools import ( file2list, - intro, - outro, format_path, - eprint, - mkOutput, + setup_logging, + mk_output, ) -from mycotools.lib.dbtools import mtdb, primaryDB -from mycotools.db2files import mtdb_main as gen_full_mtdb - -# NEED to fix when same lineage multiple ranks, e.g. Tremellales sp. will be listed -# as an order and as a genus - - -def infer_rank(db, lineage): - """Identify the taxonomic rank associated with an inputted lineage of - interest""" - linlow, rank = lineage.lower(), None - for ome, row in db.items(): - if linlow in set([x.lower() for x in row["taxonomy"].values()]): - rev_dict = { - k.lower(): v - for k, v in zip(row["taxonomy"].values(), row["taxonomy"].keys()) - } - rank = rev_dict[lineage] - - if not rank: - raise KeyError(f"no entry for {lineage}") - - return rank - - -def extract_unique(db, allowed=1, rank="species"): - """Extract unique rank from an MTDB""" - keys = copy.deepcopy(list(db.keys())) - random.shuffle(keys) - prep_db0 = {x: db[x] for x in keys} - prep_db1 = mtdb().set_index("ome") - if rank == "strain": - found = set() - for ome, row in prep_db0.items(): - name = row["taxonomy"]["species"] + " " + row["strain"] - if name not in found: - prep_db1[ome] = row - found_prep = list(found) - found_prep.append(name) - found = set(found_prep) - else: - found = defaultdict(int) - for ome, row in prep_db0.items(): - name = row["taxonomy"][rank] - found[name] += 1 - if found[name] <= allowed: - prep_db1[ome] = row - - return prep_db1 - - -def extract_tax(db, lineages): - """Extract specific taxonomic lineages of interest based on their rank""" - if isinstance(lineages, str): - lineages = [lineages] - lineages = set(x.lower() for x in lineages) - rank_dict = {k: infer_rank(db, k) for k in list(lineages)} - ranks = list(set(rank_dict.values())) - - new_db = mtdb().set_index() - for ome in db: - for rank in ranks: - try: - if db[ome]["taxonomy"][rank].lower() in lineages: - new_db[ome] = db[ome] - except KeyError: # invalid rank key for row - pass # probably should standardize tax jsons period - - return new_db - - -def extract_ome(db, omes, column="ome"): - """Extract a list of genome codes (omes) of interest""" - new_db = mtdb().set_index(column) - db = db.set_index(column) - for i in db: - if i in list(omes): - new_db[i] = db[i] - return new_db.set_index() - - -def extract_source(db, source): - """Extract an MTDB with genomes from a particular source""" - return mtdb( - { - ome: row - for ome, row in db.items() - if row["source"].lower() == source.lower() - }, - index="ome", - ) - - -def extract_pub(db): - """Extract only published and usable genomes""" - new_db = mtdb().set_index() - for ome, row in db.items(): - if row["published"]: - new_db[ome] = row - return new_db +from mycotools.lib import mtdb_sql +from mycotools.lib.dbtools import mtdb, primary_db, db_stem +from mycotools.mtdb.files import mtdb_main as gen_full_mtdb +from pathlib import Path + +logger = logging.getLogger(__name__) + + +def load_db(db_path, omes_set=(), aa_set=(), lineage_list=()): + """Load only the rows an extraction can possibly need. + + A SQLite primary database can answer "which genomes" before anything is + materialized, so an ome list, an assembly-accession list, or a set of + lineages narrows the read to an index seek. Anything else -- and any + `.mtdb` flat file -- falls back to reading the whole database.""" + if not mtdb_sql.is_sqlite(db_path): + return mtdb(db_path) + if omes_set: + return mtdb(mtdb_sql.select_omes(db_path, omes_set)) + if aa_set: + return mtdb(mtdb_sql.select_column(db_path, "assembly_acc", aa_set)) + if lineage_list: + # resolve lineages against the normalized taxonomy table; a + # species/strain lineage is not answerable there, so read everything + ranks = {mtdb_sql.infer_rank(db_path, lin) for lin in lineage_list} + if not ranks.intersection({None, "species", "strain"}): + genera = mtdb_sql.genera_for_lineages(db_path, lineage_list) + return mtdb(mtdb_sql.select_column(db_path, "genus", genera)) + return mtdb(db_path) def main( @@ -131,32 +55,33 @@ def main( nonpublished=False, inverse=False, aa_set=set(), + seed=None, ): """Python entry point for extract_mtdb""" db = db.set_index("ome") if x_number > 0: - db = extract_unique(db, x_number, rank=rank) + db = db.extract_unique(x_number, rank=rank, seed=seed) # extract each taxonomic entry based on the classification specified if lineage_list: - new_db = extract_tax(db, lineage_list) + new_db = db.extract_tax(lineage_list) # if an ome list is specified then open it, store each entry in a list and pull each ome elif omes_set: - new_db = extract_ome(db, omes_set) + new_db = db.extract_ome(omes_set) elif aa_set: - new_db = extract_ome(db, aa_set, "assembly_acc") + new_db = db.extract_ome(aa_set, "assembly_acc") # if none of these are specified then create a `new_db` variable to work for later else: new_db = db # if there is a source specified, extract it or the opposite if source: - new_db = extract_source(new_db, source) + new_db = new_db.extract_source(source) # if you want publisheds, then just pull those out if not nonpublished: - new_db = extract_pub(new_db) + new_db = new_db.extract_pub() if inverse: new_omes = set(new_db.keys()) @@ -171,9 +96,9 @@ def main( dbs = {} for lineage in lineages: if lineage: - dbs[lineage.lower()] = extract_tax(new_db, [lineage]).reset_index() + dbs[lineage.lower()] = new_db.extract_tax([lineage]).reset_index() else: - dbs["unclassified"] = extract_tax(new_db, [""]).reset_index() + dbs["unclassified"] = new_db.extract_tax([""]).reset_index() return dbs else: return new_db.reset_index() @@ -222,6 +147,11 @@ def cli(): action="store_true", help="Inverse [source|lineage(s)|nonpublished]", ) + ex_opt.add_argument( + "--seed", + type=int, + help="[-a] Seed the random sample for a reproducible selection", + ) ex_opt.add_argument("-ol", "--ome", help="File w/list of omes") ex_opt.add_argument( "-al", "--assembly_list", help="File w/list of assembly accessions" @@ -234,59 +164,53 @@ def cli(): ) out_opt.add_argument("-p", "--paths", help="Output with paths", action="store_true") out_opt.add_argument("--headers", action="store_true") - out_opt.add_argument("-d", "--mtdb", help="- for stdin", default=primaryDB()) + out_opt.add_argument("-d", "--mtdb", help="- for stdin", default=primary_db()) out_opt.add_argument("-o", "--output") args = parser.parse_args() + setup_logging(verbose=getattr(args, "verbose", False)) db_path = format_path(args.mtdb) if args.lineage or args.lineages: - eprint( - "\nWARNING: extracting taxonomy is subject to " - + "errors in NCBI's hierarchy\n" + logger.warning( + "extracting taxonomy is subject to " + "errors in NCBI's hierarchy" ) # these arguments require one another if args.by_rank and not args.rank: - eprint("\nERROR: --by_rank requires --rank", flush=True) + logger.error("--by_rank requires --rank") sys.exit(10) elif args.allowed_rank and not args.rank: - eprint("\nERROR: --allowed_rank requres --rank", flush=True) + logger.error("--allowed_rank requres --rank") sys.exit(11) elif args.rank and not args.allowed_rank and not args.by_rank: - eprint("\nERROR: --rank requires --allowed_rank or --by_rank", flush=True) + logger.error("--rank requires --allowed_rank or --by_rank") elif args.rank: args.rank = args.rank.lower() if args.rank not in set(ranks): - eprint(f"\nERROR: --rank not in {ranks}", flush=True) + logger.error(f"--rank not in {ranks}") sys.exit(12) args.lineage = args.lineage.replace('"', "").replace("'", "") output = "" if args.output: - output = format_path(args.output) - if not output.endswith("/"): - tag = "" - - if args.lineage: - tag += "_" + args.lineage - if args.lineages: - tag += "_taxonomy" - if args.source: - tag += args.source.lower() - if not args.nonpublished: - tag += "_pub" - output += "/" + os.path.basename(db_path) + tag - - if args.mtdb == "-": - data = "" - for line in sys.stdin: - data += line.rstrip() + "\n" - data = data.rstrip() - db = mtdb(data, stdin=True) - else: - db = mtdb(db_path) + # `-o` names the directory to write into; the file inside it is named + # for the source database and the filters applied. The extension is + # always `.mtdb` -- an extract is an interchange file regardless of + # which backend it was read from. + out_dir = format_path(args.output, force_dir=True) + Path(out_dir).mkdir(parents=True, exist_ok=True) + tag = "" + if args.lineage: + tag += "_" + args.lineage + if args.lineages: + tag += "_taxonomy" + if args.source: + tag += "_" + args.source.lower() + if not args.nonpublished: + tag += "_pub" + output = f"{out_dir}{db_stem(db_path)}{tag}.mtdb" if args.ome: omes = set(file2list(format_path(args.ome))) @@ -304,6 +228,15 @@ def cli(): lineage_list = [args.lineage] else: lineage_list = [] + + if args.mtdb == "-": + db = mtdb.from_string(sys.stdin.read()) + elif args.inverse: + # the inverse needs every row to subtract from, so no narrowing + db = mtdb(db_path) + else: + db = load_db(db_path, omes, aa_set, lineage_list) + new_db = main( db, lineage_list=lineage_list, @@ -315,6 +248,7 @@ def cli(): nonpublished=args.nonpublished, inverse=args.inverse, aa_set=aa_set, + seed=args.seed, ) if args.new_mtdb: gen_full_mtdb( @@ -322,13 +256,13 @@ def cli(): ) elif args.output or args.by_rank: if isinstance(new_db, mtdb): - new_db.df2db(output, paths=args.paths) + new_db.df2db(output, headers=bool(args.headers), paths=args.paths) else: - out_dir = mkOutput(output, "extract_mtdb") - prefix = re.sub(r"\.mtdb$", "", os.path.basename(db_path)) + out_dir = mk_output(output or str(Path.cwd()), "extract_mtdb") + prefix = db_stem(db_path) for lineage, db in new_db.items(): out_f = f"{out_dir}{prefix}.{lineage}.mtdb" - db.df2db(out_f) + db.df2db(out_f, headers=bool(args.headers)) else: new_db.df2db(headers=bool(args.headers), paths=args.paths) diff --git a/mycotools/db2files.py b/mycotools/mtdb/files.py similarity index 65% rename from mycotools/db2files.py rename to mycotools/mtdb/files.py index 89e9c9d..08e1665 100755 --- a/mycotools/db2files.py +++ b/mycotools/mtdb/files.py @@ -1,13 +1,17 @@ #! /usr/bin/env python3 +import logging import os -import re import sys import argparse from datetime import datetime from shutil import copy as cp -from mycotools.lib.dbtools import primaryDB, mtdb -from mycotools.lib.kontools import format_path, prep_output, eprint, vprint +from mycotools.lib import mtdb_sql +from mycotools.lib.dbtools import primary_db, mtdb +from mycotools.lib.kontools import format_path, prep_output, setup_logging +from pathlib import Path + +logger = logging.getLogger(__name__) def soft_main(filetypes, db, output_path, print_link=False, verbose=False): @@ -18,27 +22,23 @@ def soft_main(filetypes, db, output_path, print_link=False, verbose=False): if not print_link: # make the directories for each requested file type for ftype in filetypes: - if not os.path.isdir(output_path + ftype): - os.mkdir(output_path + ftype) + if not Path(output_path + ftype).is_dir(): + Path(output_path + ftype).mkdir() # grab the files for each genome code for ome, row in db.items(): for ftype in filetypes: - if os.path.isfile(row[ftype]): + if Path(row[ftype]).is_file(): sym_path = f"{output_path}{ftype}/{ome}.{ftype}" try: - os.symlink(row[ftype], sym_path) + Path(sym_path).symlink_to(row[ftype]) except FileExistsError: - if os.path.islink(sym_path): - os.remove(sym_path) - os.symlink(row[ftype], sym_path) + if Path(sym_path).is_symlink(): + Path(sym_path).unlink() + Path(sym_path).symlink_to(row[ftype]) else: - vprint( - "\t" + ome + " " + ftype + " exists", - v=verbose, - flush=True, - ) + logger.debug("" + ome + " " + ftype + " exists") else: - vprint("\tERROR: " + ome + " " + ftype, flush=True, v=verbose) + logger.debug("" + ome + " " + ftype) # simply print the link for each file else: for ome, row in db.items(): @@ -52,16 +52,16 @@ def hard_main(filetypes, db, output_path): db = db.set_index("ome") # create the directories to output each file type for ftype in filetypes: - if not os.path.isdir(output_path + ftype): - os.mkdir(output_path + ftype) + if not Path(output_path + ftype).is_dir(): + Path(output_path + ftype).mkdir() # copy each file by genome for ome, row in db.items(): for ftype in filetypes: try: - cp(row[ftype], output_path + ftype + "/" + os.path.basename(row[ftype])) + cp(row[ftype], output_path + ftype + "/" + Path(row[ftype]).name) except FileNotFoundError: - eprint("\tERROR: " + ome + " " + ftype, flush=True) + logger.error("" + ome + " " + ftype) def mtdb_main(db, output_path, og_mtdb_path): @@ -69,12 +69,12 @@ def mtdb_main(db, output_path, og_mtdb_path): # generate the base directory for output if not output_path: - output_path = os.getcwd() + "/" - if not os.path.isdir(output_path): - os.mkdir(output_path) + output_path = str(Path.cwd()) + "/" + if not Path(output_path).is_dir(): + Path(output_path).mkdir() mtdb_dir = output_path + "mycotoolsdb/" - if not os.path.isdir(mtdb_dir): - os.mkdir(mtdb_dir) + if not Path(mtdb_dir).is_dir(): + Path(mtdb_dir).mkdir() # generate the MTDB hierarchy subdirectories sub_dirs = [ @@ -84,15 +84,18 @@ def mtdb_main(db, output_path, og_mtdb_path): f"{mtdb_dir}data/", ] for dir_ in sub_dirs: - if not os.path.isdir(dir_): - os.mkdir(dir_) + if not Path(dir_).is_dir(): + Path(dir_).mkdir() # copy the og_mtdb configuration cp(og_mtdb_path + "config/mtdb.json", f"{mtdb_dir}config/mtdb.json") - # output the database + # output the database: the generated hierarchy is meant to be linked with + # `mtdb -i`, so its primary uses the SQLite backend, with a dated `.mtdb` + # snapshot alongside it for portability cdate = datetime.now().strftime("%Y%m%d") - db.df2db(f"{mtdb_dir}mtdb/{cdate}.mtdb") + db.to_sql(f"{mtdb_dir}mtdb/{mtdb_sql.PRIMARY_DB_NAME}") + db.df2db(f"{mtdb_dir}log/{cdate}.mtdb", headers=True) # output the files hard_main(["gff3", "faa", "fna"], db, f"{mtdb_dir}data/") @@ -103,7 +106,9 @@ def cli(): parser = argparse.ArgumentParser( description="Symlinks/copies selected files from database" ) - parser.add_argument("-d", "--mtdb", default=primaryDB(), help="DEFAULT: primaryDB") + parser.add_argument( + "-d", "--mtdb", default=primary_db(), help="DEFAULT: primary_db" + ) parser.add_argument("-a", "--assembly", action="store_true", help="Grab assemblies") parser.add_argument("-p", "--proteome", action="store_true", help="Grab proteomes") parser.add_argument("-g", "--gff", action="store_true", help="Grab gff`s") @@ -112,11 +117,12 @@ def cli(): parser.add_argument( "-n", "--new_mtdb", action="store_true", help="Create MTDB directory hierarchy" ) - parser.add_argument("-o", "--output", default=os.getcwd()) + parser.add_argument("-o", "--output", default=str(Path.cwd())) args = parser.parse_args() + setup_logging(verbose=getattr(args, "verbose", False)) if not args.assembly and not args.proteome and not args.gff and not args.new_mtdb: - print("\nERROR: --assembly/--proteome/--gff/--new_mtdb required", flush=True) + logger.error("--assembly/--proteome/--gff/--new_mtdb required") sys.exit(4) if args.new_mtdb: args.hard = False @@ -125,16 +131,6 @@ def cli(): db_path = format_path(args.mtdb) args.output = format_path(args.output, force_dir=True) output_path = prep_output(args.output, cd=False) - args_dict = { - "DATABASE": db_path, - "OUTPUT": output_path, - "ASSEMBLY": args.assembly, - "PROTEOME": args.proteome, - "GFF3": args.gff, - "Print links": args.print, - "Hard copy": args.hard, - "New MTDB": args.new_mtdb, - } filetypes = [] if args.proteome: diff --git a/mycotools/mtdb/manage.py b/mycotools/mtdb/manage.py new file mode 100755 index 0000000..1337486 --- /dev/null +++ b/mycotools/mtdb/manage.py @@ -0,0 +1,208 @@ +#! /usr/bin/env python3 + +import os +import sys +import logging +import argparse +from mycotools.lib.dbtools import ( + login_check, + primary_db, + mtdb, + encrypt_pw, + get_login, + store_login, +) +from mycotools.lib import mtdb_sql +from mycotools.lib.kontools import format_path, read_json, setup_logging +from pathlib import Path + +logger = logging.getLogger(__name__) + + +def ome_list(db_path): + """Every ome in the database, without materializing the rest of it.""" + if mtdb_sql.is_sqlite(db_path): + return mtdb_sql.omes(db_path) + return list(mtdb(db_path)["ome"]) + + +def rm_outdated(omes, yes=False): + """Remove outdated genomes after compiling them""" + + biofiles, to_del = [], [] + # compile the files + biofiles.extend( + [ + f"{os.environ['MYCOGFF3']}/{x}" + for x in [p.name for p in Path(os.environ["MYCOGFF3"]).iterdir()] + ] + ) + biofiles.extend( + [ + f"{os.environ['MYCOFAA']}/{x}" + for x in [p.name for p in Path(os.environ["MYCOFAA"]).iterdir()] + ] + ) + biofiles.extend( + [ + f"{os.environ['MYCOFNA']}/{x}" + for x in [p.name for p in Path(os.environ["MYCOFNA"]).iterdir()] + ] + ) + + # remove each biofile + for i in biofiles: + ome = None + ome_prep = Path(i).name + if ome_prep.endswith(".gff3"): + ome = ome_prep[:-5] + elif ome_prep.endswith(".faa"): + ome = ome_prep[:-4] + elif ome_prep.endswith(".fna"): + ome = ome_prep[:-4] + else: # safer to preserve independent placements + continue + if ome not in omes: + to_del.append(i) + + if to_del: + if yes: + data = "y" + else: + data = input(f"\n{len(to_del)} omes to be deleted.\n" + "Continue [y/N]? ") + if data.lower() in {"yes", "y"}: + for i in to_del: + Path(i).unlink() + else: + raise KeyError("cache removal stopped") + + +def restrictions( + db, restr_list, mtdb_config=format_path("~/.mycotools/config.json"), yes=False +): + mtdb_config = read_json(mtdb_config) + restr_path = mtdb_config[mtdb_config["active"]]["MYCODB"] + "../log/failed.tsv" + + try: + with open(restr_path, "r") as raw: + restricted = [x.rstrip().split() for x in raw] + except FileNotFoundError: + restricted = [] + + accs = set(x[0] for x in restricted) + for r, s, reason in restr_list: + if s.lower() in {"ncbi", "jgi"} and r not in accs: + restricted.append([r, s.lower(), str(reason)]) + logger.info("%s %s", r, s) + + in_db = [x[0] for x in restricted if x[0] in db] + while in_db: + if not yes: + check = input("Some restrictions are in the MTDB. Delete them? [y/N]: ") + if check.lower() in {"yes", "y"}: + break + else: + sys.exit(1) + + with open(restr_path, "w") as out: + out.write("\n".join(["\t".join(x) for x in restricted])) + + +def migrate(yes=False): + """Convert a flat-file primary MTDB into the SQLite backend. + + The `.mtdb` it was built from is left in place -- `primary_db()` prefers + `mtdb.db` once it exists, so the flat file becomes an inert snapshot that + can be deleted, kept for provenance, or handed to an older Mycotools.""" + db_path = primary_db() + if not db_path: + logger.error("Link a MycotoolsDB via `mtdb -i `") + return 1 + if mtdb_sql.is_sqlite(db_path): + logger.info("Primary MTDB is already SQLite: %s", db_path) + return 0 + + target = format_path("$MYCODB/" + mtdb_sql.PRIMARY_DB_NAME) + db = mtdb(db_path) + n = len(db["ome"]) + if not yes: + check = input(f"Convert {n} genomes in {db_path} to {target}? [y/N]: ") + if check.lower() not in {"yes", "y"}: + return 1 + + db.to_sql(target) + migrated = mtdb_sql.count(target) + if migrated != n: + logger.error("migrated %d of %d genomes; %s left in place", migrated, n, db_path) + return 1 + logger.info("Migrated %d genomes -> %s", migrated, target) + logger.info("%s is now a snapshot and is no longer read", db_path) + return 0 + + +def cli(): + parser = argparse.ArgumentParser( + description="Primary MycotoolsDB management utility" + ) + parser.add_argument( + "-m", + "--migrate", + action="store_true", + help="Convert a flat-file primary MTDB to the SQLite backend", + ) + parser.add_argument( + "-c", "--clear_cache", action="store_true", help="Clear MycotoolsDB legacy data" + ) + parser.add_argument( + "-p", + "--password", + action="store_true", + help="Encrypt NCBI/JGI passwords to expedite access", + ) + parser.add_argument( + "-s", + "--store", + action="store_true", + help="Store NCBI/JGI credentials WITHOUT a password (unencrypted, chmod 600)", + ) + parser.add_argument( + "-r", + "--restrict", + help="Restrict assembly accessions file, formatted: " + + "\t\t[REASON]", + ) + parser.add_argument("-y", "--yes", help="Answer yes", action="store_true") + args = parser.parse_args() + setup_logging(verbose=getattr(args, "verbose", False)) + + if args.migrate: + sys.exit(migrate(args.yes)) + + if args.password and args.store: + logger.error("--password and --store are mutually exclusive") + sys.exit(1) + if args.password: + ncbi_api, jgi_email, jgi_pwd = login_check() + encrypt_pw(ncbi_api, jgi_email, jgi_pwd) + if args.store: + ncbi_api, jgi_email, jgi_pwd = get_login(ncbi=True, jgi=True) + store_login(ncbi_api, jgi_email, jgi_pwd) + if args.restrict: + restrict_path = format_path(args.restrict) + with open(restrict_path, "r") as raw: + restricted = [x.rstrip().split("\t") for x in raw] + for v in restricted: + if len(v) < 3: + v = v + [None] + # loaded here rather than up front so the credential and migration + # operations do not pay for reading the whole database + db = mtdb(primary_db()).set_index("assembly_acc") + restrictions(db, restricted, yes=args.yes) + if args.clear_cache: + rm_outdated(set(ome_list(primary_db())), args.yes) + + sys.exit(0) + + +if __name__ == "__main__": + cli() diff --git a/mycotools/predb2mtdb.py b/mycotools/mtdb/predb.py similarity index 82% rename from mycotools/predb2mtdb.py rename to mycotools/mtdb/predb.py index 1eeeb1e..b2b5133 100755 --- a/mycotools/predb2mtdb.py +++ b/mycotools/mtdb/predb.py @@ -3,7 +3,7 @@ # NEED source to reference the annotation source # NEED to error check FAA generation simply by file size -import os +import logging import re import sys import copy @@ -11,22 +11,25 @@ import multiprocessing as mp from tqdm import tqdm from collections import Counter, defaultdict -from mycotools.lib.kontools import gunzip, mkOutput, format_path, eprint, vprint +from mycotools.lib.kontools import gunzip, mk_output, format_path from mycotools.lib.biotools import ( gff2list, list2gff, fa2dict, dict2fa, - gff3Comps, - gff2Comps, - gtfComps, + gff3_comps, + gff2_comps, + gtf_comps, ) -from mycotools.lib.dbtools import mtdb, primaryDB, loginCheck +from mycotools.lib.dbtools import mtdb, primary_db from mycotools.utils.gtf2gff3 import main as gtf2gff3 -from mycotools.utils.curGFF3 import main as curGFF3 +from mycotools.utils.cur_gff3 import main as cur_gff3 from mycotools.utils.gff2gff3 import main as gff2gff3 -from mycotools.utils.curGFF3 import rename_and_organize as rename_and_organize -from mycotools.gff2seq import aamain as gff2seq +from mycotools.utils.cur_gff3 import rename_and_organize as rename_and_organize +from mycotools.seq.gff import aamain as gff2seq +from pathlib import Path + +logger = logging.getLogger(__name__) predb_headers = [ "assembly_accession", @@ -47,7 +50,7 @@ def acq_forbid_omes(file_path): """Parse a file with forbidden ome accessions - ome codes that have been used before and are no longer valid""" - if not os.path.isfile(file_path): + if not Path(file_path).is_file(): return set() with open(file_path, "r") as raw: relics = set([x.rstrip() for x in raw]) @@ -55,12 +58,12 @@ def acq_forbid_omes(file_path): def prep_output(base_dir): - out_dir = mkOutput(base_dir, "predb2mtdb") + out_dir = mk_output(base_dir, "predb") wrk_dir = out_dir + "working/" dirs = [out_dir, wrk_dir, wrk_dir + "gff3/", wrk_dir + "fna/", wrk_dir + "faa/"] for dir_ in dirs: - if not os.path.isdir(dir_): - os.mkdir(dir_) + if not Path(dir_).is_dir(): + Path(dir_).mkdir() return dirs[:2] @@ -74,7 +77,7 @@ def copy_file(old_path, new_path): def move_biofile(old_path, ome, typ, wrk_dir, suffix=""): if old_path.endswith(".gz"): - if not os.path.isfile(old_path[:-3]): + if not Path(old_path[:-3]).is_file(): temp_path = gunzip(old_path) new_path = wrk_dir + ome + "." + typ + suffix else: @@ -83,9 +86,9 @@ def move_biofile(old_path, ome, typ, wrk_dir, suffix=""): copy_file(format_path(temp_path), new_path) else: new_path = wrk_dir + ome + "." + typ + suffix - if not os.path.isfile(new_path) and os.path.isfile(old_path): + if not Path(new_path).is_file() and Path(old_path).is_file(): copy_file(format_path(old_path), new_path) - elif not os.path.isfile(new_path): + elif not Path(new_path).is_file(): raise IOError(old_path, new_path) else: copy_file(format_path(old_path), new_path) @@ -108,14 +111,13 @@ def gen_predb(): "no", "2018", ] - eprint( + logger.info( 'INSTRUCTIONS: fill in each column with the relevant information and \ separate each column by a tab. The predb can be filled in \ via spreadsheet software and exported as a tab delimited `.tsv`. \ ASSEMBLY ACCESSIONS and PREVIOUS_OME fields must be unique to the \ - genome; otherwise predb2mtdb will update the corresponding database entry. \ - Novel data must be filled in as "new" for the genomeSource column.', - flush=True, + genome; otherwise predb will update the corresponding database entry. \ + Novel data must be filled in as "new" for the genomeSource column.' ) outputStr = "#" + "\t".join(predb_headers) outputStr += "\n#" + "\t".join(example) + "\n" @@ -183,10 +185,9 @@ def read_predb(predb_path, spacer="\t"): # required_headers.remove(head) missing_headers = required_headers.difference(set(i2header.values())) if missing_headers: - eprint( + logger.error( f"{spacer}ERROR: Required columns missing: " - + f"{missing_headers}", - flush=True, + + f"{missing_headers}" ) sys.exit(4) # if not headers: @@ -197,11 +198,8 @@ def read_predb(predb_path, spacer="\t"): # proceed with default header organization scheme if not i2header: if len(entry) != len(predb_headers): - eprint( - spacer + "ERROR: Incorrect columns, line " + str(i), - flush=True, - ) - eprint(predb_headers, "\n", entry, flush=True) + logger.error(spacer + "Incorrect columns, line " + str(i)) + logger.debug("%s\n%s", predb_headers, entry) sys.exit(3) for i1, v in enumerate(entry): predb[predb_headers[i1]].append(v.rstrip()) @@ -254,10 +252,7 @@ def read_predb(predb_path, spacer="\t"): x.lower() not in {"y", "n", "yes", "no", "", "true", "false"} for x in predb["restriction"] ): - eprint( - spacer + "ERROR: useRestriction entries must be in {y, n, yes, no}", - flush=True, - ) + logger.error(spacer + "useRestriction entries must be in {y, n, yes, no}") sys.exit(4) except KeyError: if not "restriction" in predb and "published" not in predb: @@ -266,17 +261,14 @@ def read_predb(predb_path, spacer="\t"): predb["restriction"] = [bool(x) for x in predb["published"]] if any(x.lower() not in {"jgi", "ncbi", "new"} for x in predb["source"]): - eprint( + logger.info( [ predb["assembly_acc"][i] for i, v in enumerate(predb["source"]) if v not in {"jgi", "ncbi", "new"} ] ) - eprint( - spacer + "ERROR: genomeSource entries must be in {jgi, ncbi, new}", - flush=True, - ) + logger.error(spacer + "genomeSource entries must be in {jgi, ncbi, new}") sys.exit(5) missing_from_predb = list(set(predb_headers).difference(set(predb.keys()))) @@ -384,7 +376,7 @@ def gen_omes(newdb, refdb=None, ome_col="ome", forbidden=set(), spacer="\t"): new_ome = re.sub(r"\.\d+$", "." + str(v), ome) else: new_ome = ome + ".1" # first modified version - eprint(spacer + ome + " update -> " + new_ome, flush=True) + logger.info(spacer + ome + " update -> " + new_ome) newdb["ome"][i] = new_ome continue @@ -393,18 +385,13 @@ def gen_omes(newdb, refdb=None, ome_col="ome", forbidden=set(), spacer="\t"): except TypeError: todel.append(i) if not isinstance(newdb["assembly_acc"][i], float): - eprint( - spacer - + newdb["assembly_acc"][i] - + " no metadata - " - + "failed", - flush=True, + logger.info( + spacer + newdb["assembly_acc"][i] + " no metadata - " + "failed" ) elif "index" in newdb: # for updateDB if not isinstance(newdb, float): - eprint( - spacer + newdb["index"][i] + " no metadata - " + "failed", - flush=True, + logger.info( + spacer + newdb["index"][i] + " no metadata - " + "failed" ) continue else: # no use appending failed when there's no identifiable @@ -442,7 +429,7 @@ def gen_omes(newdb, refdb=None, ome_col="ome", forbidden=set(), spacer="\t"): new_ome = re.sub(r"\.\d+$", "." + str(v), ome) else: new_ome = ome + ".1" # first modified version - eprint(spacer + ome + " update -> " + new_ome, flush=True) + logger.info(spacer + ome + " update -> " + new_ome) newdb["ome"][i] = new_ome elif ome in refdb_nover: # has a version, wasn't given in predb version_ome = refdb_nover[ome] @@ -452,7 +439,7 @@ def gen_omes(newdb, refdb=None, ome_col="ome", forbidden=set(), spacer="\t"): new_ome = ome + "." + str(v) else: raise TypeError("unknown error " + ome) - eprint(spacer + ome + " version added -> " + new_ome, flush=True) + logger.info(spacer + ome + " version added -> " + new_ome) newdb["ome"][i] = new_ome for i in reversed(todel): @@ -466,10 +453,8 @@ def cur_fna(cur_raw_fna_path, uncur_raw_fna_path, ome): ome_ver = re.search(r"(.{6}\d+).(\d+)$", ome) if ome_ver: less_ome = ome_ver[1] - ver_num = ome_ver[2] else: less_ome = ome - ver_num = 0 with open(cur_raw_fna_path + ".tmp", "w") as out: with open(uncur_raw_fna_path, "r") as in_: for line in in_: @@ -506,23 +491,20 @@ def cur_mngr( verbose=False, ): - predb_dir = os.path.basename(os.path.dirname(wrk_dir[:-1])) + "/working/" + predb_dir = Path(str(Path(wrk_dir[:-1]).parent)).name + "/working/" # assembly FNAs - vprint("\t" + ome, v=verbose, flush=True) + logger.debug("" + ome) uncur_fna_path = wrk_dir + "fna/" + ome + ".fna.uncur" cur_fna_path = wrk_dir + "fna/" + ome + ".fna" - vprint("\t\t" + predb_dir + "fna/" + ome + ".fna", v=verbose, flush=True) - if not os.path.isfile(cur_fna_path): + logger.debug("" + predb_dir + "fna/" + ome + ".fna") + if not Path(cur_fna_path).is_file(): try: uncur_fna_path = move_biofile( raw_fna_path, ome, "fa", wrk_dir + "fna/", suffix=".uncur" ) except IOError as ie: - eprint( - spacer + ome + "|" + assembly_accession + " failed FNA parsing", - flush=True, - ) + logger.info(spacer + ome + "|" + assembly_accession + " failed FNA parsing") if exit: raise ie from None return ome, False, "fna" @@ -531,16 +513,15 @@ def cur_mngr( # gene coordinate GFF3s uncur_gff_path = wrk_dir + "gff3/" + ome + ".gff3.uncur" cur_gff_path = wrk_dir + "gff3/" + ome + ".gff3" - vprint("\t\t" + predb_dir + "gff3/" + ome + ".gff3", v=verbose, flush=True) - if not os.path.isfile(cur_gff_path): + logger.debug("" + predb_dir + "gff3/" + ome + ".gff3") + if not Path(cur_gff_path).is_file(): try: uncur_gff_path = move_biofile( raw_gff_path, ome, "gff3", wrk_dir + "gff3/", suffix=".uncur" ) except IOError as ie: - eprint( - spacer + ome + "|" + assembly_accession + " failed GFF3 parsing", - flush=True, + logger.info( + spacer + ome + "|" + assembly_accession + " failed GFF3 parsing" ) if exit: raise ie from None @@ -553,9 +534,8 @@ def cur_mngr( try: gff_mngr(ome, gff, cur_gff_path, source, assembly_accession) except Exception as e: # catch all errors to continue script - eprint( - spacer + ome + "|" + assembly_accession + " failed GFF3 curation", - flush=True, + logger.info( + spacer + ome + "|" + assembly_accession + " failed GFF3 curation" ) if exit: raise e from None @@ -563,8 +543,8 @@ def cur_mngr( # proteome FAAs faa_path = wrk_dir + "faa/" + ome + ".faa" - vprint("\t\t" + predb_dir + "faa/" + ome + ".faa", v=verbose, flush=True) - if not os.path.isfile(faa_path): + logger.debug("" + predb_dir + "faa/" + ome + ".faa") + if not Path(faa_path).is_file(): try: faa = gff2seq(gff2list(cur_gff_path), fa2dict(cur_fna_path), spacer=spacer) # raise a value error if there is not a sequence for all predicted @@ -573,15 +553,13 @@ def cur_mngr( if faa and len(missing_seq) == len(faa): raise ValueError("no sequences generated in proteome") elif missing_seq: - eprint( + logger.warning( f"{spacer}\tWARNING: {len(missing_seq)} " - + "CDSs translated blank sequences", - flush=True, + + "CDSs translated blank sequences" ) except Exception as e: # catch all errors - eprint( - spacer + ome + "|" + assembly_accession + " failed proteome generation", - flush=True, + logger.info( + spacer + ome + "|" + assembly_accession + " failed proteome generation" ) if exit: raise e @@ -591,18 +569,18 @@ def cur_mngr( shutil.move(faa_path + ".tmp", faa_path) if remove: - if os.path.isfile(uncur_gff_path): - os.remove(uncur_gff_path) - if os.path.isfile(raw_gff_path): - os.remove(raw_gff_path) - if os.path.isfile(re.sub(r"\.gz$", "", raw_gff_path)): - os.remove(re.sub(r"\.gz$", "", raw_gff_path)) - if os.path.isfile(uncur_fna_path): - os.remove(uncur_fna_path) - if os.path.isfile(raw_fna_path): - os.remove(raw_fna_path) - if os.path.isfile(re.sub(r"\.gz$", "", raw_fna_path)): - os.remove(re.sub(r"\.gz$", "", raw_fna_path)) + if Path(uncur_gff_path).is_file(): + Path(uncur_gff_path).unlink() + if Path(raw_gff_path).is_file(): + Path(raw_gff_path).unlink() + if Path(re.sub(r"\.gz$", "", raw_gff_path)).is_file(): + Path(re.sub(r"\.gz$", "", raw_gff_path)).unlink() + if Path(uncur_fna_path).is_file(): + Path(uncur_fna_path).unlink() + if Path(raw_fna_path).is_file(): + Path(raw_fna_path).unlink() + if Path(re.sub(r"\.gz$", "", raw_fna_path)).is_file(): + Path(re.sub(r"\.gz$", "", raw_fna_path)).unlink() return ome, cur_fna_path, cur_gff_path, faa_path @@ -611,9 +589,9 @@ def gff_mngr(ome, gff, cur_path, source, assembly_accession): gffVer, alias = None, False for entry in gff: - if re.search(gff3Comps()["id"], entry["attributes"]): + if re.search(gff3_comps()["id"], entry["attributes"]): gffVer = 3 - alias = re.search(gff3Comps()["Alias"], entry["attributes"]) + alias = re.search(gff3_comps()["Alias"], entry["attributes"]) if alias is not None: # if entry['seqid'].startswith(ome + '_'): alias = True @@ -623,10 +601,10 @@ def gff_mngr(ome, gff, cur_path, source, assembly_accession): # entry['attributes']) else: break - elif re.search(gtfComps()["id"], entry["attributes"]): + elif re.search(gtf_comps()["id"], entry["attributes"]): gffVer = 2.5 break - elif re.search(gff2Comps()["id"], entry["attributes"]): + elif re.search(gff2_comps()["id"], entry["attributes"]): gffVer = 2 break @@ -635,25 +613,27 @@ def gff_mngr(ome, gff, cur_path, source, assembly_accession): if alias: # already curated try: new_gff = copy.deepcopy(gff) - old_ome_p = re.search(gff3Comps()["Alias"], new_gff[0]["attributes"])[1] + old_ome_p = re.search(gff3_comps()["Alias"], new_gff[0]["attributes"])[ + 1 + ] old_ome = old_ome_p[: old_ome_p.find("_")] for entry in new_gff: - # alias0 = re.search(gff3Comps()['Alias'], entry['attributes'])[1] + # alias0 = re.search(gff3_comps()['Alias'], entry['attributes'])[1] # alias_num = alias0[alias0.find('_') + 1:] # new_alias = ome + '_' + alias_num # entry['attributes'] = re.sub( - # gff3Comps()['Alias'], 'Alias='+ new_alias, + # gff3_comps()['Alias'], 'Alias='+ new_alias, # entry['attributes'] # ) entry["attributes"] = entry["attributes"].replace(old_ome, ome) new_gff = rename_and_organize(new_gff) gff = new_gff except: - gff = curGFF3(gff, ome, cur_seqids=True) + gff = cur_gff3(gff, ome, cur_seqids=True) # else: - # gff = curGFF3(gff, ome) + # gff = cur_gff3(gff, ome) else: - gff = curGFF3(gff, ome, cur_seqids=True) + gff = cur_gff3(gff, ome, cur_seqids=True) elif gffVer == 2.5: gff, trans_str, failed, flagged = gtf2gff3(gff, ome) else: @@ -662,7 +642,6 @@ def gff_mngr(ome, gff, cur_path, source, assembly_accession): ver_search = re.search(r"(.{6}\d+)\.(\d+)", ome) if ver_search is not None: less_ome = ver_search[1] - ome_ver = ver_search[2] else: less_ome = ome for line in gff: @@ -699,11 +678,11 @@ def main( ): for dir_ in [wrk_dir + "fna/", wrk_dir + "gff3/", wrk_dir + "faa/"]: - if not os.path.isdir(dir_): - os.mkdir(dir_) + if not Path(dir_).is_dir(): + Path(dir_).mkdir() infdb = predb2mtdb(predb) - vprint("\nGenerating omes", v=verbose, flush=True) + logger.debug("Generating omes") omedb, failed = gen_omes( infdb, refdb, ome_col="ome", forbidden=forbidden, spacer=spacer ) @@ -726,7 +705,7 @@ def main( ] ) - vprint("\nCurating data", v=verbose, flush=True) + logger.debug("Curating data") if cpus > 1: with mp.Pool(processes=cpus) as pool: cur_data = pool.starmap(cur_mngr, tqdm(cur_cmds, total=len(cur_cmds))) @@ -750,14 +729,14 @@ def main( def cli(): usage = ( - "Generate a predb file:\npredb2mtdb\n\nCreate a mycotoolsdb " - + "from a predb file:\npredb2mtdb \n\nCreate a mycotoolsdb " - + "referencing an alternative master database:\npredb2mtdb " - + "\nSkip failing genomes:\npredb2mtdb -s" + "Generate a predb file:\nmtdb predb\n\nCreate a mycotoolsdb " + + "from a predb file:\nmtdb predb \n\nCreate a mycotoolsdb " + + "referencing an alternative master database:\nmtdb predb " + + "\nSkip failing genomes:\nmtdb predb -s" ) if any(x in {"-h", "--help", "-help"} for x in sys.argv): - eprint("\n" + usage + "\n", flush=True) + logger.info("" + usage + "\n") sys.exit(0) elif len(sys.argv) == 1: print(gen_predb()) @@ -768,9 +747,9 @@ def cli(): elif len(sys.argv) > 3: refDB = mtdb(format_path(sys.argv[3])) else: - refDB = mtdb(primaryDB()) + refDB = mtdb(primary_db()) else: - refDB = mtdb(primaryDB()) + refDB = mtdb(primary_db()) if set(sys.argv).intersection({"-s", "--skip"}): exit = False @@ -778,14 +757,14 @@ def cli(): exit = True # from Bio import Entrez - # ncbi_email, ncbi_api, jgi_email, jgi_pwd = loginCheck(jgi = False) + # ncbi_email, ncbi_api, jgi_email, jgi_pwd = login_check(jgi = False) # Entrez.email = ncbi_email # if ncbi_api: # Entrez.api_key = ncbi_api - eprint("\nPreparing run", flush=True) + logger.info("Preparing run") predb = read_predb(format_path(sys.argv[1]), spacer="\t") - out_dir, wrk_dir = prep_output(os.path.dirname(format_path(sys.argv[1]))) + out_dir, wrk_dir = prep_output(str(Path(format_path(sys.argv[1])).parent)) forbid_omes = acq_forbid_omes(file_path=format_path("$MYCODB/../log/relics.txt")) @@ -796,8 +775,8 @@ def cli(): # from mycotools.lib.dbtools import gather_taxonomy, assimilate_tax # tax_dicts = gather_taxonomy(omedb, api_key = ncbi_api) # outdb, genus_dicts = assimilate_tax(omedb, tax_dicts) - # outdb.df2db(out_dir + 'predb2mtdb.mtdb') - omedb.df2db(out_dir + "predb2mtdb.mtdb") + # outdb.df2db(out_dir + 'predb.mtdb') + omedb.df2db(out_dir + "predb.mtdb") sys.exit(0) diff --git a/mycotools/update_mtdb.py b/mycotools/mtdb/update.py similarity index 58% rename from mycotools/update_mtdb.py rename to mycotools/mtdb/update.py index f1e72fa..93dbd7c 100755 --- a/mycotools/update_mtdb.py +++ b/mycotools/mtdb/update.py @@ -2,7 +2,6 @@ # NEED reinit implementation # NEED to update introduction -# NEED a verbose option # NEED revert version option (ome-by-ome/list of omes) # NEED to reference a manually curated duplicate check # NEED a prohibit option to import prohibited JGI/NCBI IDs and option to update @@ -12,18 +11,17 @@ # will require logging whatever NCBI omes directly overlap MycoCosm # NEED to remove overlap when rerunning failed genomes +import logging import os import re import sys -import time import json -import base64 +import time import shutil -import getpass -import hashlib import zipfile import requests import argparse +import warnings import subprocess import numpy as np import pandas as pd @@ -31,13 +29,13 @@ from Bio import Entrez from datetime import datetime from collections import defaultdict +from mycotools.lib import mtdb_sql from mycotools.lib.dbtools import ( db2df, df2db, gather_taxonomy, - assimilate_tax, - primaryDB, - loginCheck, + primary_db, + login_check, log_editor, mtdb, mtdb_initialize, @@ -46,29 +44,49 @@ intro, outro, format_path, - eprint, prep_output, collect_files, read_json, write_json, split_input, - findExecs, + find_execs, + setup_logging, + atomic_write, ) from mycotools.lib.biotools import fa2dict, gff2list, dict2fa, list2gff -from mycotools.ncbiDwnld import ( +from mycotools.download.ncbi import ( esearch_ncbi, esummary_ncbi, run_datasets, compile_organism_names, main as ncbiDwnld, ) -from mycotools.jgiDwnld import main as jgiDwnld +from mycotools.download.jgi import main as jgiDwnld from mycotools.utils.ncbi2db import main as ncbi2db from mycotools.utils.jgi2db import main as jgi2db -from mycotools.predb2mtdb import main as predb2mtdb -from mycotools.predb2mtdb import predb_headers, read_predb, gen_omes -from mycotools.assemblyStats import main as assStats -from mycotools.annotationStats import main as annStats +from mycotools.mtdb.predb import main as predb2mtdb +from mycotools.mtdb.predb import predb_headers, read_predb, gen_omes +from pathlib import Path + +logger = logging.getLogger(__name__) + + +def _read_ledger(file_path, comment="#"): + """Return the meaningful lines of a ledger/sidecar file. + + Yields the shared read half of the ledger parsers below: returns a list of + right-stripped lines, skipping comment lines (those beginning with + `comment`) and lines that are empty once stripped. Returns an empty list if + the file does not exist, so callers can treat a missing ledger as empty. + """ + if not Path(file_path).is_file(): + return [] + with open(file_path, "r") as raw: + return [ + line.rstrip() + for line in raw + if not line.startswith(comment) and line.strip() + ] def validate_t_and_c(config, discrepancy=False): @@ -89,14 +107,13 @@ def validate_t_and_c(config, discrepancy=False): # if there isnt a configuration, alert the user to use-restriction policies else: - print( - "\nPlease review JGI use-restricted data policy here: " + logger.debug( + "Please review JGI use-restricted data policy here: " + "https://jgi.doe.gov/user-programs/pmo-overview/policies/" + "\nPlease review GenBank use-restricted data policy here: " + "https://ncbi.nlm.nih.gov/genbank/" + "\nPlease review how Mycotools handles use-restricted data here:" - + " https://github.com/xonq/mycotools/blob/master/MTDB.md", - flush=True, + + " https://github.com/xonq/mycotools/blob/master/MTDB.md" ) check = "" if check.lower() not in {"y", "yes"}: @@ -112,7 +129,7 @@ def validate_t_and_c(config, discrepancy=False): + "\n\nPlease type [y]es/[N]o if you acknowledge these terms: " ) if check.lower() not in {"y", "yes"}: - print("\nRerun without --nonpublished", flush=True) + logger.info("Rerun without --nonpublished") sys.exit(1) nonpublished = "yes" @@ -123,8 +140,6 @@ def validate_t_and_c(config, discrepancy=False): def gen_config( branch="fungi", forbidden="", - repo=None, - rogue=False, nonpublished=False, jgi=False, rank2lineages={}, @@ -132,10 +147,8 @@ def gen_config( config = { "forbidden": forbidden, - "repository": repo, "branch": branch, "nonpublished": nonpublished, - "rogue": rogue, "jgi": jgi, "lineage_constraints": rank2lineages, } @@ -143,25 +156,21 @@ def gen_config( return config -def add_vars(init_dir, dbtype): - """Initialize the environmental variables for MTDB""" - mtdb_initialize(init_dir, dbtype, init=True) - - -def initDB( - init_dir, - branch, - envs, - dbtype, - date=None, - rogue=False, - nonpublished=False, - jgi=True, - repo=None, - rank2lineages={}, -): +def init_db(init_dir): """Initialize database in `init_dir`""" + init_dir = format_path(init) + if Path(init_dir).is_dir(): + init_dir += "mycotoolsdb/" + if not init_dir.endswith("/"): + init_dir += "/" + envs = { + "MYCOFNA": init_dir + "data/fna", + "MYCOFAA": init_dir + "data/faa", + "MYCOGFF3": init_dir + "data/gff3", + "MYCODB": init_dir + "mtdb/", + } + os.environ["MYCODB"] = init_dir + "mtdb/" new_dirs = [ init_dir + "data/", init_dir + "config/", @@ -176,184 +185,99 @@ def initDB( if not output.endswith("/"): output += "/" for new_dir in new_dirs: - if not os.path.isdir(new_dir): - os.mkdir(new_dir) - - config = gen_config( - branch=branch, - rogue=rogue, - forbidden="$MYCODB/log/forbidden.tsv", - nonpublished=nonpublished, - jgi=jgi, - repo=repo, - rank2lineages=rank2lineages, - ) - write_json(config, init_dir + "config/mtdb.json", indent=1) - - if not rogue: - # this is a relic, and needs to be adjusted to a central reference if - # that is ever created - if not os.path.isdir(init_dir + "mtdb"): - # NEED TO CHANGE FROM SSH TO LINK ONCE OPEN (config['repository']) - git_exit = subprocess.call( - [ - "git", - "clone", - "git@gitlab.com:xonq/mtdb", - init_dir + "mtdb", - # '-b', branch - ] - ) - if git_exit != 0: - eprint("\nERROR: git clone failed.", flush=True) - sys.exit(2) - else: - print("\nmycotoolsdb directory already exists", flush=True) - # NEED TO ADD GITIGNORE TO GIT - if not primaryDB(): - eprint( - "\nERROR: no YYYYmmdd.mtdb in " + format_path(envs["MYCODB"]), - flush=True, - ) - sys.exit(3) - else: - new_db_path = output + "mtdb/" + date + ".mtdb" - if not os.path.isfile(new_db_path): - with open(output + "mtdb/" + date + ".mtdb", "w") as out: - out.write("".join(["\t" for x in mtdb.columns])) + if not Path(new_dir).is_dir(): + Path(new_dir).mkdir() - return output, config + new_db_path = output + "mtdb/" + mtdb_sql.PRIMARY_DB_NAME + if not Path(new_db_path).is_file(): + mtdb().to_sql(new_db_path) + for env in envs: + os.environ[env] = envs[env] + init_db = db2df(mtdb()) # initialize a new database + mtdb_initialize( + init_dir, init=True + ) + return init_db, f"{output}log/{date}/" -def parse_forbidden(forbidden_path): - """Read the log file containing information on forbidden genomes""" - # NEED to be rewritten to reference a central repository forbidden file - - log_path = format_path(forbidden_path) - if os.path.isfile(log_path): - log_dict = readLog(log_path) - else: - log_dict = {} - - return log_dict - - -def add_forbidden(tag, source, file_path=None, flag="failed download"): - edit = tag + "\t" + source + "\t" + flag - log_editor(file_path, tag, edit) def parse_dups(file_path): - """Retrieve a file containing replicated genomes and ignore these. - This is important for dereplication of discrepant genus naming between NCBI - and MycoCosm due to not adhering to conserved genus naming standards and - updating relic genus names. It appears that MycoCosm will name genera by - their anamorph occassionally, and not the consensus name - though I assume - this is also to some extent present in NCBI. Ultimately, a manually curated - file is necessary for this and should be held in a central repository.""" + """Parse a file of replicated genomes to ignore, for dereplicating + discrepant NCBI/MycoCosm genus naming (e.g. MycoCosm sometimes names a genus + by its anamorph rather than the consensus). Requires a manually curated + file, ideally kept in a central repository.""" duplicates = {} - if os.path.isfile(file_path): - with open(file_path, "r") as raw: - for line in raw: - if not line.startswith("#"): - data = [x.rstrip() for x in line.split("\t") if x] - if data: - duplicates[data[0]] = [data[1], data[2], data[3]] - + for line in _read_ledger(file_path): + data = [x.rstrip() for x in line.split("\t") if x] + if data: + duplicates[data[0]] = [data[1], data[2], data[3]] return duplicates -# def add_dups( -# dup_code, dup_entry, file_path -# ): -# edit = dup_code + '\t' + '\t'.join(dup_entry) -# log_editor(file_path, dup_code, edit) - - def acq_forbid_omes(file_path): """Parse a file with forbidden ome accessions - ome codes that have been used before and are no longer valid""" - if not os.path.isfile(file_path): - return set() - with open(file_path, "r") as raw: - relics = set([x.rstrip() for x in raw]) - return relics + return set(_read_ledger(file_path)) def write_forbid_omes(omes, file_path): """Add to a file of forbidden ome accessions so that these are not ever used again, even if the codename is removed from the database""" - if os.path.isfile(file_path): - with open(file_path, "r") as raw: - old_relics = set([x.rstrip() for x in raw]) - new_relics = old_relics.union(set(omes)) - else: - new_relics = set(omes) - - with open(file_path + ".tmp", "w") as out: # be cautious because if it - # cancels then we lose the old data + # union with any pre-existing relics; atomic_write guards against losing + # the old data if the write is cancelled midway + new_relics = set(_read_ledger(file_path)).union(set(omes)) + with atomic_write(file_path) as out: out.write("\n".join([str(x) for x in sorted(new_relics)])) - shutil.move(file_path + ".tmp", file_path) def parse_failed(file_path=None, rerun=False): """Parse a file that stores the failed accessions and metadata of the attempted acquisition. Return a dictionary that contains the failed accession and its metadata.""" - prev_failed = {} - if not os.path.isfile(file_path) or rerun: + if not Path(file_path).is_file() or rerun: with open(file_path, "w") as out: out.write("#code\tsource\tversion\tattempt_date") - else: - with open(file_path, "r") as raw: - for line in raw: - if not line.startswith("#"): - data = [x.rstrip() for x in line.split("\t")] - while len(data) < 4: - data.append("") - prev_failed[data[0]] = { - "source": data[1], - "version": data[2], - "attempt_date": data[3], - } - + return {} + prev_failed = {} + for line in _read_ledger(file_path): + data = [x.rstrip() for x in line.split("\t")] + while len(data) < 4: + data.append("") + prev_failed[data[0]] = { + "source": data[1], + "version": data[2], + "attempt_date": data[3], + } return prev_failed def parse_jgi2ncbi(file_path): """Parse previously collected NCBI to JGI data to limit querying""" - jgi2ncbi = {} - if not os.path.isfile(file_path): + if not Path(file_path).is_file(): with open(file_path, "w") as out: out.write("#ncbi_acc\tmycocosm_portal") - else: - with open(file_path, "r") as raw: - for line in raw: - if not line.startswith("#"): - d = line.rstrip().split("\t") - ncbi, jgi = d[0], d[1].lower() - jgi2ncbi[jgi] = ncbi - + return {} + jgi2ncbi = {} + for line in _read_ledger(file_path): + d = line.split("\t") + jgi2ncbi[d[1].lower()] = d[0] return jgi2ncbi def parse_true_ncbi(file_path): """Parse accessions considered to be unique to NCBI""" - true_ncbi = set() - if not os.path.isfile(file_path): + if not Path(file_path).is_file(): with open(file_path, "w") as out: out.write("#ncbi_acc") - else: - with open(file_path, "r") as raw: - true_ncbi = set([x.rstrip() for x in raw if not x.startswith("#")]) - - return true_ncbi + return set() + return set(_read_ledger(file_path)) def add_true_ncbi(true_ncbi, file_path=None): """Add to a ledger of accessions considered to be unique to NCBI""" - with open(file_path, "w") as out: - out.write("#ncbi_acc\n" + "\n".join([str(x) for x in list(true_ncbi)])) + with atomic_write(file_path) as out: + out.write("#ncbi_acc\n" + "\n".join([str(x) for x in true_ncbi])) def add_jgi2ncbi(jgi2ncbi, file_path=None): @@ -361,7 +285,7 @@ def add_jgi2ncbi(jgi2ncbi, file_path=None): is prone to failure given that the field JGI uses to supply their genome accession is either absent from some NCBI entries, or is in a different field""" - with open(file_path, "w") as out: + with atomic_write(file_path) as out: out.write("#ncbi_acc\tmycocosm_portal\n") for jgi, ncbi in jgi2ncbi.items(): out.write(ncbi + "\t" + jgi + "\n") @@ -377,44 +301,110 @@ def add_failed(code, source, version, date, file_path): log_editor(file_path, code, edit) +def read_mycocosm(table_path): + """Parse a downloaded MycoCosm table into a dataframe, or raise ValueError + if the file is not that table. + + Encodings are attempted strictest-first. MycoCosm does not declare one, and + the table carries non-ASCII characters in its publication and strain + fields, but single-byte codecs such as cp1252 and latin1 map nearly every + byte, so they do not raise on input that is not theirs -- they silently + decode it to mojibake. Only utf-8 rejects bytes that are not its own, so + leading with it is what makes the fallback capable of choosing at all.""" + + decode_error = None + for encoding in ("utf-8", "cp1252", "latin1"): + try: + jgi_df = pd.read_csv(table_path, encoding=encoding) + except UnicodeDecodeError as error: + decode_error = error + continue + except pd.errors.ParserError as error: + # not a decoding problem: the bytes were text, but not a table + raise ValueError(f"{table_path} is not parseable as CSV: {error}") + + jgi_df.columns = [x.replace('"', "") for x in jgi_df.columns] + missing = {"name", "portal"}.difference(set(jgi_df.columns)) + if missing: + raise ValueError( + f"{table_path} lacks the MycoCosm column(s) " + + f"{'/'.join(sorted(missing))}" + ) + return jgi_df + + raise ValueError(f"{table_path} could not be decoded: {decode_error}") + + def dwnld_mycocosm( out_file, mycocosm_url="https://mycocosm.jgi.doe.gov/ext-api/mycocosm/catalog/" + "download-group?flt=&seq=all&pub=all&grp=fungi&srt=" + "released&ord=desc", + max_attempts=3, ): - """Download the MycoCosm genome data spreadsheet, format to UTF-8 and - return a Pandas dataframe of the data""" + """Download the MycoCosm genome data spreadsheet and return a Pandas + dataframe of the data. - check_curl = findExecs(["curl"], verbose=False) + The download is only accepted once it parses as the MycoCosm table. JGI + answers a failed request with a 404 HTML page, which curl reports as a + success unless it is asked not to, so an outage would otherwise be cached + to out_file as though it were data -- and because the presence of out_file + is what marks the download done, every later run would reread the error + page rather than retry the download.""" - if not os.path.isfile(out_file): - for attempt in range(3): + check_curl = find_execs(["curl"], verbose=False) + + if Path(out_file).is_file(): + try: + return read_mycocosm(out_file) + except ValueError as error: + logger.warning(f"discarding unusable MycoCosm table: {error}") + Path(out_file).unlink() + + for attempt in range(1, max_attempts + 1): + if attempt > 1: + time.sleep(3 * 2 ** (attempt - 2)) + try: if check_curl: - curl_cmd = subprocess.call( - ["curl", mycocosm_url, "-o", out_file + ".tmp"], + # --fail turns an HTTP error into a non-zero exit rather than a + # saved error page; --location follows the redirect JGI issues. + # curl's stderr is captured rather than inherited so that its + # 404 does not print over the run -- the exit status is what + # this needs, and the message is only worth a debug line + subprocess.run( + ["curl", "--fail", "--location", "--silent", "--show-error", + mycocosm_url, "-o", out_file + ".tmp"], + check=True, stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, ) - if not curl_cmd: - shutil.move(out_file + ".tmp", out_file) - break - if curl_cmd: - eprint("\nERROR: failed to retrieve MycoCosm table", flush=True) else: - resp = requests.get(url) + resp = requests.get(mycocosm_url) + resp.raise_for_status() with open(out_file + ".tmp", "wb") as f: f.write(resp.content) - shutil.move(out_file + ".tmp", out_file) - - try: - jgi_df = pd.read_csv(out_file, encoding="cp1252") - except UnicodeDecodeError: - jgi_df = pd.read_csv(out_file, encoding="latin1") - except UnicodeDecodeError: - jgi_df = pd.read_csv(out_file, encoding="utf-8") - jgi_df.columns = [x.replace('"', "").replace('"', "") for x in jgi_df.columns] - - return jgi_df + jgi_df = read_mycocosm(out_file + ".tmp") + except (subprocess.CalledProcessError, requests.RequestException, + ValueError) as error: + # one failed attempt is not news until they have all failed, and + # during a JGI outage this fires every time -- reporting each 404 + # buries the run in noise that says nothing the final message does + # not say once + logger.debug(f"MycoCosm table attempt {attempt} failed: {error}") + Path(out_file + ".tmp").unlink(missing_ok=True) + continue + shutil.move(out_file + ".tmp", out_file) + return jgi_df + + # every attempt failed the same way, which is far more often JGI being down + # than anything wrong locally -- MycoCosm's ext-api has served nothing but + # its 404 page for days at a time, and it is what this URL resolves to + logger.error( + "MycoCosm may be inaccessible; no table was retrieved in " + + f"{max_attempts} attempts" + ) + sys.exit(24) def dwnld_ncbi_metadata( @@ -426,130 +416,180 @@ def dwnld_ncbi_metadata( eukaryotes, and return a Pandas dataframe""" ncbi_url = ncbi_url + group + ".txt" - if not os.path.isfile(ncbi_file): - getTbl = subprocess.call(["curl", ncbi_url, "-o", ncbi_file + ".tmp"]) + if not Path(ncbi_file).is_file(): + subprocess.call(["curl", ncbi_url, "-o", ncbi_file + ".tmp"]) shutil.move(ncbi_file + ".tmp", ncbi_file) ncbi_df = pd.read_csv(ncbi_file, sep="\t") return ncbi_df -def prep_taxa_cols( - df, taxonomy_dir, col="#Organism/Name", api=None, acc2org={}, max_attempts=3 +def dwnld_data_reports( + accs, + out_dir, + api=None, + report_chunk=1000, + max_attempts=3, + exit_code=11, + label="GenBank", ): - - skip_prep = list(acc2org.keys()) - gca_prep = [x.upper().replace("GCF", "GCA") for x in skip_prep] - gcf_prep = [x.upper().replace("GCA", "GCF") for x in skip_prep] - skip = set(gca_prep + gcf_prep) - if not os.path.isdir(taxonomy_dir): - os.mkdir(taxonomy_dir) - aa_file = taxonomy_dir + "assembly_accs.genbank.txt" - with open(aa_file, "w") as out: - out.write("\n".join([x for x in list(df["assembly_acc"]) if x not in skip])) - - attempts, datasets_cmd = 0, 0 - if not os.path.isdir(taxonomy_dir + "ncbi_dataset"): - datasets_path = taxonomy_dir + "ncbi_dataset.zip" - while attempts < max_attempts: - if attempts: - eprint("\t\t\tReattempting", flush=True) - if os.path.isfile(datasets_path): - os.remove(datasets_path) - attempts += 1 - datasets_cmd = run_datasets( - None, aa_file, taxonomy_dir, True, api=api, verbose=True - ) - try: - with zipfile.ZipFile(datasets_path, "r") as zip_ref: - zip_ref.extractall(taxonomy_dir) - os.remove(datasets_path) - break - except zipfile.BadZipFile: - eprint( - f"\t\tERROR: datasets download corrupted - {attempts}", flush=True - ) - if attempts == max_attempts: - sys.exit(11) - except FileNotFoundError: - eprint(f"\t\tERROR: datasets failed - {attempts}", flush=True) - - if datasets_cmd: - eprint(f"\t\tWARNING: datasets failed, assuming no genomes found", flush=True) - acc2org_n, acc2meta = {}, {} - else: - acc2org_n, acc2meta, org_failed = compile_organism_names( - taxonomy_dir + "ncbi_dataset/" - ) - print( - f"\t\t{len(acc2meta) + len(org_failed)}", - "genomes queried from GenBank", - flush=True, + """Acquire NCBI assembly data reports; return acc2org, acc2meta, failed. + + `datasets` is called on `report_chunk` accessions at a time rather than on + the whole list at once. A from-scratch initialization queries >16,000 + accessions, and one request that large is both slow enough to be dropped + mid-transfer and all-or-nothing when it is: a single failure discards every + record. Each chunk downloads into its own directory and is left there, so an + interrupted run resumes at the first chunk it had not finished. + + These are dehydrated metadata records rather than genomes, so a chunk is a + few hundred KB and is sized well above the genome chunk: the transfer is + nowhere near large enough to be what NCBI resets, and a smaller size would + only multiply the requests a from-scratch run makes.""" + acc2org, acc2meta, failed = {}, {}, [] + empty_chunks = [] + + # a run started before chunking landed leaves a single whole-list download + # here; parse it rather than re-acquiring everything + legacy_dir = out_dir + "ncbi_dataset/" + if Path(legacy_dir + "data/assembly_data_report.jsonl").is_file(): + return compile_organism_names(legacy_dir) + + acc_chunks = [ + accs[i : i + report_chunk] for i in range(0, len(accs), report_chunk) + ] + # `datasets` draws its own per-call bar, which is meaningless here because + # it restarts every chunk; it is silenced below in favor of one bar over the + # whole acquisition + for chunk_i, acc_chunk in enumerate( + tqdm( + acc_chunks, + total=len(acc_chunks), + desc=f"{label} reports", + unit=" chunk", + disable=len(acc_chunks) < 2, ) - print(f'\t\t{len(org_failed)/len(df["assembly_acc"])*100}% failed', flush=True) - - # check for RefSeq for failed entries - refseq_dir = taxonomy_dir + "refseq/" - if not os.path.isdir(refseq_dir): - os.mkdir(refseq_dir) - missing_accs = sorted( - set(df["assembly_acc"]).difference( - set(acc2org_n.keys()).union(set(acc2org.keys())) + ): + chunk_dir = f"{out_dir}chunk_{chunk_i}/" + acc_file = f"{chunk_dir}assembly_accs.txt" + unzip_dir = f"{chunk_dir}ncbi_dataset/" + zip_path = f"{chunk_dir}ncbi_dataset.zip" + if not Path(chunk_dir).is_dir(): + Path(chunk_dir).mkdir(parents=True) + + # only reuse a completed chunk that covers exactly these accessions; + # otherwise the chunk boundaries have shifted since it was written and + # its contents no longer correspond to this index + expected = "\n".join(acc_chunk) + cached = ( + Path(unzip_dir).is_dir() + and Path(acc_file).is_file() + and Path(acc_file).read_text() == expected ) - ) - reattempt_acc = [] - for acc in missing_accs: - if acc.upper().startswith("GCA"): - reattempt_acc.append(acc.upper().replace("GCA_", "GCF_")) - elif acc.upper().startswith("GCF"): - reattempt_acc.append(acc.upper().replace("GCF_", "GCA_")) - acc_file_re = refseq_dir + "assembly_accs.refseq.txt" - with open(acc_file_re, "w") as out: - out.write("\n".join(reattempt_acc)) - - # attempt to download the ncbi_datasets zip file until allowed attempts are - # exhausted - rs_datasets_cmd = 0 - if not os.path.isdir(refseq_dir + "ncbi_dataset"): - print(f"\t\tChecking RefSeq for {len(reattempt_acc)} entries", flush=True) - rs_datasets_path = refseq_dir + "ncbi_dataset.zip" - attempts = 0 - while attempts < max_attempts: - if attempts: - eprint("\t\t\t\tReattempting", flush=True) - if os.path.isfile(rs_datasets_path): - os.remove(rs_datasets_path) - attempts += 1 - rs_datasets_cmd = run_datasets( - None, acc_file_re, refseq_dir, True, api=api, verbose=True - ) - try: - with zipfile.ZipFile(rs_datasets_path, "r") as zip_ref: - zip_ref.extractall(refseq_dir) - os.remove(rs_datasets_path) - break - except zipfile.BadZipFile: - eprint( - f"\t\t\tERROR: datasets download corrupted - {attempts}", flush=True + if not cached: + shutil.rmtree(unzip_dir, ignore_errors=True) + with open(acc_file, "w") as out: + out.write(expected) + # every message in this loop is debug: a retry is routine, it would + # overdraw the bar above, and a chunk that never succeeds is + # accounted for in the single summary once the bar has finished + attempts = 0 + while attempts < max_attempts: + if attempts: + logger.debug(f"Reattempting {label} chunk {chunk_i + 1}") + if Path(zip_path).is_file(): + Path(zip_path).unlink() + attempts += 1 + run_datasets( + None, + acc_file, + chunk_dir, + True, + api=api, + verbose=False, + # datasets' own stderr would overdraw the chunk bar above; + # it is kept at debug, and the failures below are what the + # user is told about + mute_stderr=True, ) - if attempts == max_attempts: - sys.exit(10) - except FileNotFoundError: - eprint(f"\t\tERROR: datasets failed - {attempts}", flush=True) - - if rs_datasets_cmd: - eprint(f"\t\tWARNING: datasets failed, assuming no genomes found", flush=True) - acc2org_rs, acc2meta_rs = {}, {} - else: - acc2org_rs, acc2meta_rs, org_failed_2 = compile_organism_names( - refseq_dir + "ncbi_dataset/" + try: + with zipfile.ZipFile(zip_path, "r") as zip_ref: + zip_ref.extractall(chunk_dir) + Path(zip_path).unlink() + break + except zipfile.BadZipFile: + # a truncated archive is corruption rather than absence, and + # exhausting the attempts on it is fatal -- so this one is + # said out loud + if attempts == max_attempts: + logger.error(f"{label} chunk {chunk_i + 1} download corrupted") + sys.exit(exit_code) + logger.debug(f"datasets download corrupted - {attempts}") + except FileNotFoundError: + logger.debug(f"datasets failed - {attempts}") + + # datasets exited without writing a report on every attempt: the chunk + # holds no retrievable genomes, which is not fatal to the remainder + if not Path(unzip_dir + "data/assembly_data_report.jsonl").is_file(): + empty_chunks.append(chunk_i + 1) + continue + + c_acc2org, c_acc2meta, c_failed = compile_organism_names(unzip_dir) + acc2org.update(c_acc2org) + acc2meta.update(c_acc2meta) + failed.extend(c_failed) + logger.debug( + f"\t\t{label} chunk {chunk_i + 1}/{len(acc_chunks)}: " + + f"{len(c_acc2meta)} genomes" ) - print(f"\t\t{len(acc2meta_rs)} genome(s) queried from RefSeq", flush=True) - acc2org, acc2meta = {**acc2org, **acc2org_n, **acc2org_rs}, { - **acc2meta, - **acc2meta_rs, - } + # one line once the bar is done, rather than a burst of them through it; + # empty chunks are routine for the RefSeq pass, where most accessions are + # speculative, so this reports the scale and leaves the detail to -v + if empty_chunks: + logger.warning( + f"{len(empty_chunks)}/{len(acc_chunks)} {label} chunk(s) returned " + + f"no genomes ({len(empty_chunks) * report_chunk} accessions at " + + "most); " + + "rerun with -v for the datasets output" + ) + logger.debug(f"empty {label} chunks: {empty_chunks}") + + return acc2org, acc2meta, failed + + +def prep_taxa_cols( + df, + taxonomy_dir, + col="#Organism/Name", + api=None, + acc2org={}, + max_attempts=3, + report_chunk=1000, +): + + skip_prep = list(acc2org.keys()) + skip = set(x.upper().replace("GCF", "GCA") for x in skip_prep) + if not Path(taxonomy_dir).is_dir(): + Path(taxonomy_dir).mkdir() + + gb_accs = [x for x in list(df["assembly_acc"]) if x not in skip] + acc2org_n, acc2meta, org_failed = dwnld_data_reports( + gb_accs, + taxonomy_dir, + api=api, + report_chunk=report_chunk, + max_attempts=max_attempts, + exit_code=11, + label="GenBank", + ) + logger.info( + "%s %s", + f"\t\t{len(acc2meta) + len(org_failed)}", + "genomes queried from GenBank", + ) + + acc2org = {**acc2org, **acc2org_n} df["strain"] = "" todel = set() @@ -597,12 +637,15 @@ def prep_jgi_cols(jgi_df, name_col="name"): return jgi_df -def clean_ncbi_df(ncbi_df, update_path, kingdom="Fungi", api=None, max_attempts=3): - ncbi_df = ncbi_df.astype(str).replace(np.nan, "") +def clean_ncbi_df( + ncbi_df_init, update_path, kingdom="Fungi", api=None, max_attempts=3, + report_chunk=1000 +): + ncbi_df = ncbi_df_init.rename(columns={"Assembly Accession": "assembly_acc"}).astype(str).replace(np.nan, "") acc2org_path = update_path + "../gca2org.tsv" acc2org = {} - if os.path.isfile(acc2org_path): + if Path(acc2org_path).is_file(): with open(acc2org_path, "r") as raw: for line in raw: d = line.split("\t") @@ -626,14 +669,17 @@ def clean_ncbi_df(ncbi_df, update_path, kingdom="Fungi", api=None, max_attempts= ncbi_df = ncbi_df[ncbi_df["assembly_acc"].str.startswith(("GCA", "GCF"))] ncbi_df, acc2meta, acc2org = prep_taxa_cols( - ncbi_df, update_path + "taxonomy/", api=api, acc2org=acc2org + ncbi_df, + update_path + "taxonomy/", + api=api, + acc2org=acc2org, + report_chunk=report_chunk, ) - with open(acc2org_path + ".tmp", "w") as out: + with atomic_write(acc2org_path) as out: for acc, org in acc2org.items(): org_meta = f'{org["genus"]}\t{org["species"]}\t{org["strain"]}' out.write(f"{acc}\t{org_meta}\n") - os.rename(acc2org_path + ".tmp", acc2org_path) # remove entries without sufficient metadata ncbi_df = ncbi_df.dropna(subset=["genus"]) @@ -750,21 +796,46 @@ def rm_ncbi_overlap(ncbi_df, mycocosm_df, jgi2ncbi, fails=set(), acc2meta={}, ap jgi2biosample[gen_sp] = row["BioSample Accession"] else: fails.add(row["assembly_acc"]) - # for i in reversed(todel): - # ncbi_jgi_overlap = pd.concat([ncbi_jgi_overlap, ncbi_df.loc[i]]) - # ncbi_df = ncbi_df.drop(i) ncbi_df, ncbi_jgi_overlap = exec_rm_overlap(ncbi_df, todel) return ncbi_df, jgi2ncbi, jgi2biosample, fails, ncbi_jgi_overlap, todel +def write_primary(db, date, update_path=None): + """Install `db` as the primary MTDB. + + The primary is the SQLite database at `$MYCODB/mtdb.db`, written atomically + so an interrupted update can never leave a partial database where + `primary_db()` would pick it up. The outgoing primary is archived under + `log//` first, alongside a `.mtdb` snapshot of the new one -- every + historical primary stays readable with nothing but a text editor.""" + new_path = format_path("$MYCODB/" + mtdb_sql.PRIMARY_DB_NAME) + prior = primary_db(verbose=False) + + # copy, rather than move, so a failed write leaves the old primary in place + if update_path and prior and Path(prior).is_file(): + archive = update_path + Path(prior).name + if format_path(prior) != format_path(archive): + shutil.copy(prior, archive) + + db.to_sql(new_path) + + if update_path: + db.df2db(update_path + date + ".mtdb", headers=True) + # a dated flat primary predates the SQLite backend; it has been archived, so + # drop it rather than leave a stale database beside the real one + if prior and Path(prior).is_file() and format_path(prior) != format_path(new_path): + Path(prior).unlink() + return new_path + + def mk_wrk_dirs(update_path): """Make the download directories in the update path""" wrk_dirs = ["faa/", "fna/", "gff3/"] for wrk_dir in wrk_dirs: - if not os.path.isdir(update_path + wrk_dir): - os.mkdir(update_path + wrk_dir) + if not Path(update_path + wrk_dir).is_dir(): + Path(update_path + wrk_dir).mkdir() def prepare_ref_db(ref_db, date): @@ -780,7 +851,7 @@ def prepare_ref_db(ref_db, date): {k: v for k, v in ref_db.items() if v["source"].lower() == "ncbi"}, index="ome" ) if set(ref_db.keys()).difference(set(jgi.keys()).union(set(ncbi.keys()))): - eprint( + logger.warning( '\tWARNING: reference entries that are not labeled "jgi/ncbi" are excluded' ) @@ -852,11 +923,9 @@ def internal_redundancy_check(db): def read_prev_tax(tax_path): """Open a genus to taxonomy JSON path""" tax_dicts = {} - if os.path.isfile(tax_path): - with open(tax_path, "r") as raw: - for line in raw: - data = line.rstrip().split("\t") - tax_dicts[data[0]] = json.loads(data[1]) + for line in _read_ledger(tax_path): + data = line.split("\t") + tax_dicts[data[0]] = json.loads(data[1]) return tax_dicts @@ -868,35 +937,42 @@ def ref_update( jgi_email, jgi_pwd, config, - ncbi_email, ncbi_api, cpus=1, check_MD5=True, - jgi=True, group="eukaryotes", kingdom="Fungi", remove=True, taxonomy=True, - ncbi_fallback=False, + chunk=25, + tape_wait=None, ): """Initialize/Update the primary MTDB based on a reference database acquired external from any existing primary MTDB""" # NEED to mark none for new databases' refdb # initialize update - print("\nInitializing run", flush=True) + logger.info("Initializing run") mk_wrk_dirs(update_path) jgi_df, ncbi_df = prepare_ref_db(ref_db, date) # run JGI - if jgi and len(jgi_df) > 0: - print("\nAssimilating MycoCosm", flush=True) - jgi_db_path = update_path + date + ".jgi.mtdb" + if jgi_email and len(jgi_df) > 0: + logger.info("Assimilating MycoCosm") jgi_predb_path = update_path + date + ".jgi.predb2.mtdb" - if not os.path.isfile(jgi_predb_path): - print("\tDownloading MycoCosm data", flush=True) - post_jgi_df, jgi_failed = jgiDwnld(jgi_df, update_path, jgi_email, jgi_pwd) + if not Path(jgi_predb_path).is_file(): + logger.info("Downloading MycoCosm data") + jgi_deferred = set() + post_jgi_df, jgi_dwnld_failed = jgiDwnld( + jgi_df, + update_path, + jgi_email, + jgi_pwd, + deferred=jgi_deferred, + restore_wait=tape_wait, + defer_tape=True, + ) jgi_predb = post_jgi_df.rename( columns={ "published(s)": "published", @@ -905,19 +981,41 @@ def ref_update( } ) - print("\tCurating MycoCosm data", flush=True) + logger.info("Curating MycoCosm data") jgi_premtdb = jgi_predb.fillna("").to_dict(orient="list") - jgi_mtdb, jgi_failed1 = predb2mtdb( - jgi_premtdb, - mtdb(), - update_path, - # forbidden = forbid_omes, - cpus=cpus, - remove=remove, - spacer="\t\t", - ) - jgi_failed = list(jgi_failed) - jgi_failed.extend(jgi_failed1) + # jgiDwnld reports bare portal ids; the failed ledger records + # [accession, version] pairs. Genomes only awaiting a JGI tape + # restore are retryable, so they are not recorded as failures + versions = dict(zip(jgi_df["assembly_acc"], jgi_df["version"])) + jgi_failed = [ + [acc, versions.get(acc, "")] + for acc in jgi_dwnld_failed + if acc not in jgi_deferred + ] + if jgi_deferred: + logger.info( + f"\t{len(jgi_deferred)} genome(s) awaiting JGI tape restore; " + "they will be retried on the next run" + ) + # a downloaded assembly path is required to curate; if no JGI genome + # was successfully retrieved (e.g. all portals failed) skip curation + # rather than raising a KeyError and aborting the whole run + if "assemblyPath" in jgi_premtdb: + jgi_mtdb, jgi_failed1 = predb2mtdb( + jgi_premtdb, + mtdb(), + update_path, + # forbidden = forbid_omes, + cpus=cpus, + remove=remove, + spacer="\t\t", + ) + jgi_failed.extend(jgi_failed1) + else: + logger.warning( + "No JGI assemblies downloaded; skipping MycoCosm curation" + ) + jgi_mtdb = mtdb() jgi_mtdb.df2db(jgi_predb_path) for failure in jgi_failed: add_failed( @@ -935,38 +1033,22 @@ def ref_update( jgi_mtdb = mtdb() new_db = jgi_mtdb.mtdb2pd() - print("\nAssimilating NCBI", flush=True) - if not os.path.isfile(update_path + date + ".ncbi.predb"): - print("\tDownloading NCBI data", flush=True) - if ncbi_fallback: - from mycotools.ncbi_dwnld_fallback import main as ncbi_dwnld_fallback - - ncbi_predb, ncbi_failed1 = ncbi_dwnld_fallback( - assembly=True, - proteome=False, - gff3=True, - ncbi_df=ncbi_df, - remove=True, - output_path=update_path, - column="assembly_acc", - ncbi_column="genome", - check_MD5=check_MD5, - verbose=True, - ) - - else: - ncbi_predb, ncbi_failed1 = ncbiDwnld( - assembly=True, - proteome=False, - gff3=True, - ncbi_df=ncbi_df, - remove=True, - output_path=update_path, - column="assembly_acc", - ncbi_column="genome", - check_MD5=check_MD5, - verbose=True, - ) + logger.info("Assimilating NCBI") + if not Path(update_path + date + ".ncbi.predb").is_file(): + logger.info("Downloading NCBI data") + ncbi_predb, ncbi_failed1 = ncbiDwnld( + assembly=True, + proteome=False, + gff3=True, + ncbi_df=ncbi_df, + remove=True, + output_path=update_path, + column="assembly_acc", + ncbi_column="genome", + check_MD5=check_MD5, + verbose=True, + chunk=chunk, + ) for failure in ncbi_failed1: add_failed( @@ -984,8 +1066,8 @@ def ref_update( # refdbncbi = mtdb(update_path + date + '.ncbi.ref.mtdb') ncbi_predb = pd.read_csv(update_path + date + ".ncbi.predb", sep="\t") - print("\tCurating NCBI data", flush=True) - if not os.path.isfile(update_path + date + ".ncbi.predb2.mtdb"): + logger.info("Curating NCBI data") + if not Path(update_path + date + ".ncbi.predb2.mtdb").is_file(): for key in ncbi_predb.columns: ncbi_predb[key] = ncbi_predb[key].fillna("") ncbi_predb["version"] = ncbi_predb["version"].astype(str) @@ -1019,7 +1101,7 @@ def ref_update( # df2db(ncbi_db, ncbi_db_path) new_db = pd.concat([new_db, ncbi_db]) - print("\nAssimilating NCBI taxonomy data", flush=True) + logger.info("Assimilating NCBI taxonomy data") new_mtdb = mtdb.pd2mtdb(new_db) if kingdom.lower() == "fungi": @@ -1036,7 +1118,7 @@ def ref_update( elif jgi_mtdb: update_mtdb = jgi_mtdb else: - eprint("\nNo updates", flush=True) + logger.info("No updates") sys.exit(0) if taxonomy: # already completed @@ -1050,8 +1132,7 @@ def ref_update( output_path=tax_path, tax_dicts=tax_dicts, ) - new_mtdb, genus_dicts = assimilate_tax(new_mtdb, tax_dicts) - dupFiles = {"fna": {}, "faa": {}, "gff3": {}} + genus_dicts = new_mtdb.assimilate_tax(tax_dicts) for ome, row in update_mtdb.items(): if row["genus"] in genus_dicts: @@ -1107,37 +1188,34 @@ def extract_constraint_lineages( passing_tax.add(genus) df = df[df["genus"].isin(passing_tax)] + return tax_dicts, df def taxonomy_update( - orig_db, + db, update_path, date, config, - ncbi_email, ncbi_api, rank="kingdom", group="fungi", ): """Reset the taxonomy for the entire database and overwrite the previous tax path data to accomodate new taxonomy""" - taxless_db = orig_db.reset_index() - taxless_db["taxonomy"] = [{} for x in taxless_db["taxonomy"]] + db = db.reset_index() + db["taxonomy"] = [{} for x in db["taxonomy"]] tax_path = f"{update_path}../taxonomy.tsv" gca_path = f"{update_path}../gca2org.tsv" - if os.path.isfile(tax_path): - os.rename(tax_path, update_path + "old_taxonomy.tsv") - if os.path.isfile(gca_path): - os.rename(gca_path, update_path + "old_gca2org.tsv") + if Path(tax_path).is_file(): + Path(tax_path).rename(update_path + "old_taxonomy.tsv") + if Path(gca_path).is_file(): + Path(gca_path).rename(update_path + "old_gca2org.tsv") tax_dicts = gather_taxonomy( - taxless_db, api_key=ncbi_api, king=group, rank=rank, output_path=tax_path + db, api_key=ncbi_api, king=group, rank=rank, output_path=tax_path ) - tax_db, genus_dicts = assimilate_tax(taxless_db, tax_dicts) - if not isinstance(tax_db, mtdb): - return tax_db, mtdb.pd2mtdb(tax_db) - else: - return tax_db.mtdb2pd(), tax_db + genus_dicts = db.assimilate_tax(tax_dicts) + return db def rogue_update( @@ -1148,21 +1226,20 @@ def rogue_update( jgi_email, jgi_pwd, config, - ncbi_email, ncbi_api, cpus=1, check_MD5=True, - jgi=True, group="eukaryotes", kingdom="Fungi", remove=True, lineage_constraints={}, - ncbi_fallback=False, + chunk=25, + report_chunk=1000, + tape_wait=None, ): """Initialize/update a standalone primary MTDB""" - # NEED to mark none for new databases' refdb # initialize update - print("\nInitializing run", flush=True) + logger.info("Initializing run") mk_wrk_dirs(update_path) prev_failed = parse_failed( rerun=rerun, file_path=format_path("$MYCODB/../log/failed.tsv") @@ -1178,23 +1255,23 @@ def rogue_update( else: api = 3 ncbi_db_path = update_path + date + ".ncbi.mtdb" - pre_ncbi_df0 = dwnld_ncbi_metadata(update_path + date + ".ncbi.tsv", group=group) - pre_ncbi_df1 = pre_ncbi_df0.rename(columns={"Assembly Accession": "assembly_acc"}) - print("\tAcquiring NCBI metadata", flush=True) + ncbi_df_init = dwnld_ncbi_metadata(update_path + date + ".ncbi.tsv", group=group) + logger.info("Acquiring NCBI metadata") ncbi_df, acc2meta = clean_ncbi_df( - pre_ncbi_df1, update_path, kingdom=kingdom, api=ncbi_api + ncbi_df_init, + update_path, + kingdom=kingdom, + api=ncbi_api, + report_chunk=report_chunk, ) # begin extracting lineages of interest and store tax_dicts for later tax_path = f"{update_path}../taxonomy.tsv" tax_dicts = read_prev_tax(tax_path) - # tax_dicts = {v['genus']: v['taxonomy'] for k, v in db.iterrows() \ - # if any(y for x, y in v['taxonomy'].items() \ - # if x not in {'genus', 'species', 'strain'})} if lineage_constraints: lineage_path = update_path + date + ".ncbi.posttax.df" - if not os.path.isfile(lineage_path): - print("\nExtracting lineages from NCBI", flush=True) + if not Path(lineage_path).is_file(): + logger.info("Extracting lineages from NCBI") # NEED to transition to datasets tax_dicts, ncbi_df = extract_constraint_lineages( ncbi_df, ncbi_api, kingdom, lineage_constraints, tax_dicts, tax_path @@ -1207,11 +1284,11 @@ def rogue_update( old_len = len(db["ome"]) new_len = len(db["ome"]) if old_len - new_len: - print("\t" + str(old_len - new_len) + " redundant entries removed", flush=True) + logger.debug("" + str(old_len - new_len) + " redundant entries removed") # run JGI - if jgi: - print("\nAssimilating MycoCosm (1 download/minute)", flush=True) + if jgi_email: + logger.info("Assimilating MycoCosm (1 download/minute)") jgi_db_path = update_path + date + ".jgi.mtdb" mycocosm_path = update_path + date + ".mycocosm.csv" @@ -1223,8 +1300,8 @@ def rogue_update( # extract JGI lineages of interest and store tax_dicts for later if lineage_constraints: lineage_path = update_path + date + ".jgi.posttax.df" - if not os.path.isfile(lineage_path): - print("\tExtracting lineages from MycoCosm", flush=True) + if not Path(lineage_path).is_file(): + logger.info("Extracting lineages from MycoCosm") tax_dicts, jgi_df = extract_constraint_lineages( jgi_df, ncbi_api, kingdom, lineage_constraints, tax_dicts, tax_path ) @@ -1232,11 +1309,11 @@ def rogue_update( else: jgi_df = pd.read_csv(lineage_path, sep="\t") - print("\tSearching NCBI for MycoCosm overlap", flush=True) + logger.info("Searching NCBI for MycoCosm overlap") jgi_ncbi_overlap_file = f"{update_path}/redundant_ncbi.tsv" jgi2ncbi = parse_jgi2ncbi(update_path + "../jgi2ncbi.tsv") ncbi_df = ncbi_df.set_index("assembly_acc", drop=False) - if os.path.isfile(jgi_ncbi_overlap_file): + if Path(jgi_ncbi_overlap_file).is_file(): with open(jgi_ncbi_overlap_file, "r") as raw: todel_i = [x.rstrip() for x in raw] ncbi_df, ncbi_jgi_overlap = exec_rm_overlap(ncbi_df, todel_i) @@ -1246,10 +1323,9 @@ def rogue_update( rm_ncbi_overlap(ncbi_df, jgi_df, jgi2ncbi, true_ncbi, acc2meta, api=api) ) - print("\t\t" + str(len(jgi2ncbi)) + " overlapping genomes", flush=True) - with open(jgi_ncbi_overlap_file + ".tmp", "w") as out: + logger.debug("" + str(len(jgi2ncbi)) + " overlapping genomes") + with atomic_write(jgi_ncbi_overlap_file) as out: out.write("\n".join([x for x in todel_i])) - os.rename(jgi_ncbi_overlap_file + ".tmp", jgi_ncbi_overlap_file) add_true_ncbi(true_ncbi, update_path + "../supported_ncbi.tsv") add_jgi2ncbi(jgi2ncbi, update_path + "../jgi2ncbi.tsv") for i, row in jgi_df.iterrows(): @@ -1259,9 +1335,9 @@ def rogue_update( ): jgi_df.at[i, "biosample"] = jgi2biosample[row["portal"].lower()] - print("\tDownloading MycoCosm data", flush=True) + logger.info("Downloading MycoCosm data") jgi_predb_path = update_path + date + ".jgi.predb2.mtdb" - jgi_predb, db, jgi_failed = jgi2db( + jgi_predb, db, jgi_failed, jgi_deferred = jgi2db( jgi_df, db, update_path, @@ -1273,18 +1349,23 @@ def rogue_update( failed_dict=prev_failed, jgi2ncbi=jgi2ncbi, repeatmasked=True, + restore_wait=tape_wait, ) # download JGI files and ready predb - # get ncbi hits that hit failed jgi runs - failed_ncbi2jgi = {jgi2ncbi[f[0]]: f[0] for f in jgi_failed if f[0] in jgi2ncbi} + # get ncbi hits that hit jgi runs which did not yield data - genomes + # awaiting a JGI tape restore included, so NCBI covers them until the + # next run retrieves the MycoCosm copy + failed_ncbi2jgi = { + jgi2ncbi[f[0]]: f[0] for f in jgi_failed + jgi_deferred if f[0] in jgi2ncbi + } ncbi_jgi_overlap = ncbi_jgi_overlap[ ncbi_jgi_overlap["assembly_acc"].isin(failed_ncbi2jgi) ] ncbi_df = pd.concat([ncbi_df, ncbi_jgi_overlap]) refdbjgi = mtdb.pd2mtdb(db) - if not os.path.isfile(jgi_predb_path): - print("\tCurating MycoCosm data", flush=True) + if not Path(jgi_predb_path).is_file(): + logger.info("Curating MycoCosm data") jgi_premtdb = jgi_predb.fillna("").to_dict(orient="list") if "assemblyPath" in jgi_premtdb: jgi_mtdb, jgi_failed1 = predb2mtdb( @@ -1301,6 +1382,8 @@ def rogue_update( jgi_mtdb = mtdb() jgi_mtdb.df2db(jgi_predb_path) + # only genuine failures are blacklisted; genomes pending a JGI tape + # restore are absent from jgi_failed so the next run retries them for failure in jgi_failed: add_failed( failure[0], @@ -1318,7 +1401,7 @@ def rogue_update( jgi_db = pd.DataFrame({x: [] for x in refdbjgi.keys()}) new_db_path = update_path + date + ".checkpoint.jgi.mtdb" - if not os.path.isfile(new_db_path): + if not Path(new_db_path).is_file(): if len(jgi_db) > 0: df2db(jgi_db, jgi_db_path) if not db is None: @@ -1333,13 +1416,12 @@ def rogue_update( else: jgi_mtdb = None new_db = db - new_dups = duplicates - print("\nAssimilating NCBI (10 download/second w/API key, 3 w/o)", flush=True) + logger.info("Assimilating NCBI (10 download/second w/API key, 3 w/o)") new_db["version"] = new_db["version"].astype(str) - if not os.path.isfile(update_path + date + ".ncbi.predb"): + if not Path(update_path + date + ".ncbi.predb").is_file(): # if not os.path.isfile(update_path + date + '.ncbi.predb'): - print("\tDownloading NCBI data", flush=True) + logger.info("Downloading NCBI data") ncbi_predb, new_db, ncbi_failed1 = ncbi2db( update_path, ncbi_df, @@ -1349,7 +1431,7 @@ def rogue_update( rerun=rerun, duplicates=duplicates, check_MD5=check_MD5, - fallback=ncbi_fallback, + chunk=chunk, ) for failure in ncbi_failed1: @@ -1360,8 +1442,6 @@ def rogue_update( date, format_path("$MYCODB/../log/failed.tsv"), ) - # for dup in new_dups: - # add_dups(dup, new_dups[dup], format_path('$MYCODB/../log/duplicates.tsv')) refdbncbi = mtdb.pd2mtdb(new_db) refdbncbi.df2db(update_path + date + ".ncbi.ref.mtdb") ncbi_predb.to_csv(update_path + date + ".ncbi.predb", sep="\t", index=None) @@ -1369,8 +1449,8 @@ def rogue_update( refdbncbi = mtdb(update_path + date + ".ncbi.ref.mtdb") ncbi_predb = pd.read_csv(update_path + date + ".ncbi.predb", sep="\t") - print("\tCurating NCBI data", flush=True) - if not os.path.isfile(update_path + date + ".ncbi.predb2.mtdb"): + logger.info("Curating NCBI data") + if not Path(update_path + date + ".ncbi.predb2.mtdb").is_file(): for key in ncbi_predb.columns: ncbi_predb[key] = ncbi_predb[key].fillna("") ncbi_predb["version"] = ncbi_predb["version"].astype(str) @@ -1410,7 +1490,7 @@ def rogue_update( new_mtdb = mtdb.pd2mtdb(new_db) - print("\nAssimilating NCBI taxonomy data", flush=True) + logger.info("Assimilating NCBI taxonomy data") if kingdom.lower() == "fungi": rank = "kingdom" else: @@ -1425,8 +1505,6 @@ def rogue_update( tax_dicts=tax_dicts, output_path=tax_path, ) - new_mtdb, genus_dicts = assimilate_tax(new_mtdb, tax_dicts) - dupFiles = {"fna": {}, "faa": {}, "gff3": {}} if jgi_mtdb and ncbi_mtdb: update_mtdb = mtdb( @@ -1437,56 +1515,22 @@ def rogue_update( elif jgi_mtdb: update_mtdb = jgi_mtdb else: - eprint("\nNo updates", flush=True) + logger.info("No updates") sys.exit(0) + genus_dicts = new_mtdb.assimilate_tax(tax_dicts) + update_mtdb.assimilate_tax(genus_dicts) + return new_mtdb, update_mtdb def rm_raw_data(out_dir): """Remove raw data after completion""" for i in ["faa", "gff3", "gff", "xml", "fna"]: - if os.path.isdir(out_dir + i): + if Path(out_dir + i).is_dir(): shutil.rmtree(out_dir + i) -def gen_algn_db(update_path, omes): - """Generate an alignment database for the complete primary MTDB""" - date = os.path.basename(os.path.abspath(update_path)) - fas = collect_files(os.environ["MYCOFAA"] + "/", ".faa") - fas = [x for x in fas if os.path.basename(x)[:-6] in omes] - mkdb_base = "cat " + " ".join(fas) - mkdb_blast = ( - mkdb_base - + " | makeblastdb -in -" - + " -out " - + os.environ["MYCOGFF3"] - + "../db/" - + date - + ".db -parse_seqids -dbtype prot -title " - + date - + ".db" - ) - # mkdb_mmseqs = mkdb_base + ' | mmseqs createdb stdin ' + \ - # format_path('$MYCOFAA/' + date + '.mmseqs.db') + '; ' + \ - # 'mmseqs createdb ' + format_path('$MYCOFAA/' + date + \ - # '.mmseqs.db') + ' tmp' - with open(update_path + date + "_makeblastdb.sh", "w") as out: - out.write(mkdb_blast) - # with open(update_path + date + '_mmseqsdb.sh', 'w') as out: - # out.write(mkdb_mmseqs) - - print( - "\nOPTIONAL: To generate blastdb | mmseqsdb, run the following" - + "\nbash " - + update_path - + date - + "_makeblastdb.sh" - ) - # bash ' + update_path \ - # + date + '_mmseqsdb.sh') - - def check_add_mtdb(orig_mtdb, add_mtdb, update_path, overwrite=True): """Check the original MTDB for overlapping omes and curate if necessary""" orig_mtdb = orig_mtdb.set_index() @@ -1510,10 +1554,6 @@ def check_add_mtdb(orig_mtdb, add_mtdb, update_path, overwrite=True): else: for ome in overwrite_omes: del add_mtdb[ome] - # eprint('\nERROR: assembly accessions ("assembly_acc") must be ' \ - # 'unique between databases: ', flush = True) - # eprint(', '.join(failed_aas), flush = True) - # sys.exit(123) orig_omes = set(orig_mtdb.keys()) new_omes = set(add_mtdb.keys()) @@ -1552,7 +1592,7 @@ def check_add_mtdb(orig_mtdb, add_mtdb, update_path, overwrite=True): new_ome2old_ome = {v: k for k, v in old_ome2new_ome.items()} new_ome_mtdb = new_ome_mtdb.set_index("ome") for k, v in old_ome2new_ome.items(): - print(f"\t{k} converted to {v}", flush=True) + logger.debug(f"{k} converted to {v}") # create directories for new files fna_dir, gff_dir, faa_dir = ( @@ -1561,8 +1601,8 @@ def check_add_mtdb(orig_mtdb, add_mtdb, update_path, overwrite=True): f"{update_path}faa/", ) for path_ in [fna_dir, gff_dir, faa_dir]: - if not os.path.isdir(path_): - os.mkdir(path_) + if not Path(path_).is_dir(): + Path(path_).mkdir() # convert the file header names to the new omes for ome, old_ome in new_ome2old_ome.items(): @@ -1599,7 +1639,7 @@ def check_add_mtdb(orig_mtdb, add_mtdb, update_path, overwrite=True): return add_mtdb.reset_index() -def db2primary(addDB, refDB, save=False, combined=False): +def db2primary(add_mtdb, refDB, save=False, combined=False): """Finalize an update by converting the updated MTDB into the primary MTDB""" if save: @@ -1607,71 +1647,80 @@ def db2primary(addDB, refDB, save=False, combined=False): else: move_ns = shutil.move - addDB = addDB.reset_index() + add_mtdb = add_mtdb.reset_index() refDB = refDB.reset_index() refOmes = set(refDB["ome"]) - addOmes = set(addDB["ome"]) + addOmes = set(add_mtdb["ome"]) base_ome2update_ome = {re.search(r"^[^\d]+\d+", x)[0]: x for x in refDB["ome"] if x} updates = {} refDB = refDB.set_index() if refOmes.intersection(addOmes) and not combined: - eprint(refOmes.intersection(addOmes), flush=True) + logger.info(refOmes.intersection(addOmes)) raise KeyError( - "ERROR: ome codes exist in database. Rerun predb2mtdb or remove manually" + "ERROR: ome codes exist in database. Rerun `mtdb predb` or remove manually" ) - for i, ome in enumerate(addDB["ome"]): + for i, ome in enumerate(add_mtdb["ome"]): base_ome = re.search(r"^[^\d]+\d+", ome)[0] if base_ome in base_ome2update_ome: update_ome = base_ome2update_ome[base_ome] updates[update_ome] = ome del refDB[update_ome] - if os.path.isfile(addDB["gff3"][i]): - move_ns(addDB["gff3"][i], format_path("$MYCOGFF3/" + ome + ".gff3")) - elif not os.path.isfile(format_path("$MYCOGFF3/" + ome + ".gff3")): + if Path(add_mtdb["gff3"][i]).is_file(): + move_ns(add_mtdb["gff3"][i], format_path("$MYCOGFF3/" + ome + ".gff3")) + elif not Path(format_path("$MYCOGFF3/" + ome + ".gff3")).is_file(): raise FileNotFoundError(f"{ome} missing gff3 for unknown reason") - if os.path.isfile(addDB["fna"][i]): - move_ns(addDB["fna"][i], format_path("$MYCOFNA/" + ome + ".fna")) - elif not os.path.isfile(format_path("$MYCOFNA/" + ome + ".fna")): + if Path(add_mtdb["fna"][i]).is_file(): + move_ns(add_mtdb["fna"][i], format_path("$MYCOFNA/" + ome + ".fna")) + elif not Path(format_path("$MYCOFNA/" + ome + ".fna")).is_file(): raise FileNotFoundError(f"{ome} missing fna for unknown reason") - if os.path.isfile(addDB["faa"][i]): - move_ns(addDB["faa"][i], format_path("$MYCOFAA/" + ome + ".faa")) - elif not os.path.isfile(format_path("$MYCOFAA/" + ome + ".faa")): + if Path(add_mtdb["faa"][i]).is_file(): + move_ns(add_mtdb["faa"][i], format_path("$MYCOFAA/" + ome + ".faa")) + elif not Path(format_path("$MYCOFAA/" + ome + ".faa")).is_file(): raise FileNotFoundError(f"{ome} missing faa for unknown reason") - addDB["gff3"][i] = os.environ["MYCOGFF3"] + ome + ".gff3" - addDB["fna"][i] = os.environ["MYCOFNA"] + ome + ".fna" - addDB["faa"][i] = os.environ["MYCOFAA"] + ome + ".faa" - addDB = addDB.set_index() - for ome, row in addDB.items(): + add_mtdb["gff3"][i] = os.environ["MYCOGFF3"] + ome + ".gff3" + add_mtdb["fna"][i] = os.environ["MYCOFNA"] + ome + ".fna" + add_mtdb["faa"][i] = os.environ["MYCOFAA"] + ome + ".faa" + add_mtdb = add_mtdb.set_index() + for ome, row in add_mtdb.items(): refDB[ome] = row return refDB.reset_index(), updates -def control_flow( - init, - update, - reference, - add, - taxonomy, - predb, - save, - nonpublished, - ncbi_only, - lineage, - rank, - kingdom, - failed, - forbidden, - resume, - no_md5, - cpu, - ncbi_email=False, - ncbi_api=None, - overwrite=True, - fallback=False, -): +def error_handle_args(args): + if not any(x for x in [args.init, args.update, args.reference, args.add, args.taxonomy]): + raise ValueError("--update/--init/--reference/--add/--taxonomy must be specified") + elif args.reference and not args.init: + raise ValueError("--reference requires a --init directory") + elif args.lineage and not args.rank: + raise ValueError("--lineage requires --rank") + elif args.lineage and not args.init: + raise ValueError("--lineage requires --init") + elif args.predb and not args.init: + raise ValueError("--predb requires --init") + elif args.predb and args.lineage: + raise ValueError("--predb and --lineage are incompatible") + # persistent configuration is fixed at initialization; after a database + # exists it is changed through `mtdb configure`, not a plain update + elif args.nonpublished and not args.init: + raise ValueError( + "--nonpublished is set at initialization; " + "change it afterward via `mtdb configure --nonpublished`" + ) + elif args.ncbi_only and not args.init: + raise ValueError( + "--ncbi_only is set at initialization; " + "change it afterward via `mtdb configure --ncbi_only`" + ) + elif args.reference: + if args.add: + raise ValueError("--add and --reference are incompatible") + elif args.predb: + raise ValueError("--reference and --predb are incompatible") + +def determine_kingdom(raw_kingdom): abbr2king = { "a": "animals", "r": "archaea", @@ -1680,54 +1729,16 @@ def control_flow( "p": "plants", } - kingdom = kingdom.lower() + kingdom = raw_kingdom.lower() if kingdom not in abbr2king: if kingdom not in set(abbr2king.values()): - eprint("\nERROR: invalid --kingdom", flush=True) - sys.exit(431) + raise KeyError("invalid --kingdom") else: kingdom = abbr2king[kingdom] + return kingdom - if not init and not update and not reference and not add and not taxonomy: - eprint( - "\nERROR: --update/--init/--reference/--add must be specified", flush=True - ) - sys.exit(15) - elif reference and not init: - eprint("\nERROR: --reference requires a --init directory", flush=True) - sys.exit(14) - elif lineage and not rank: - eprint("\nERROR: --lineage requires --rank") - sys.exit(16) - elif lineage and not init: - eprint("\nERROR: --lineage requires --init") - sys.exit(17) - elif predb and not init: - eprint("\nERROR: --predb requires --init") - sys.exit(18) - elif predb and lineage: - eprint("\nERROR: --predb and --lineage are incompatible") - sys.exit(20) - elif reference: - if add: - eprint("\nERROR: --add and --reference are incompatible") - sys.exit(13) - elif predb: - eprint("\nERROR: --reference and --predb are incompatible") - sys.exit(19) - else: - ref_db = mtdb(format_path(reference), add_paths=False) - - if predb: - predb_path = format_path(predb) - - # if rogue: - rogue_bool = True - if ncbi_only: - jgi = False - else: - jgi = True +def parse_lineages(lineage, rank): # acquire the lineages inputted rank2lineages = {} permitted_ranks = {"phylum", "subphylum", "class", "order", "family", "genus"} @@ -1735,12 +1746,10 @@ def control_flow( lineage_constraints = split_input(lineage) rank_constraints = split_input(rank) if len(lineage_constraints) != len(rank_constraints): - eprint("\nERROR: --lineage must be same length as --rank") - sys.exit(18) + raise ValueError("--lineage must be same length as --rank") for rank_c in rank_constraints: if rank_c.lower() not in permitted_ranks: - eprint(f"\nERROR: accepted ranks: {permitted_ranks}") - sys.exit(22) + raise KeyError(f"accepted ranks: {permitted_ranks}") rank2lineages = defaultdict(set) for i, v in enumerate(lineage_constraints): rank2lineages[rank_constraints[i]].add(v.lower()) @@ -1748,334 +1757,178 @@ def control_flow( k.lower(): sorted(v) for k, v in sorted(rank2lineages.items(), key=lambda x: x[0]) } + return rank2lineages + + +def parse_nonpublished(kingdom, nonpublished, config): + # nonfungi is nonpublished by default because it is all GenBank + if kingdom != "fungi": + return True + elif nonpublished: + return validate_t_and_c(config) + else: + return False + +def import_config(init, kingdom, nonpublished, jgi, rank2lineages, config_path=format_path("$MYCODB/../config/mtdb.json")): # parse and check configuration nonpublished arguments config = {} if "MYCODB" in os.environ: - config_path = format_path("$MYCODB/../config/mtdb.json") - if os.path.isfile(config_path): + if Path(config_path).is_file(): config = read_json(format_path(config_path)) # for LEGACY installs: if "lineage_constraints" not in config: config["lineage_constraints"] = {} write_json(config, config_path) elif not init: - eprint("\nERROR: corrupted MycotoolsDB - no configuration found") - sys.exit(21) - if not init: # is MYCODB initialized? - # rogue_bool = config['rogue'] - # nonpublished = config['nonpublished'] - if bool(nonpublished) and not bool(config["nonpublished"]): - config["nonpublished"] = validate_t_and_c(config, discrepancy=True) - write_json(config, config_path) - if bool(config["jgi"]) and bool(ncbi_only): # and not overwrite: - eprint( - "\nERROR: --ncbi_only specified after initialization", flush=True - ) - sys.exit(173) - elif init: + raise FileNotFoundError(f"no configuration found at {config_path}") + # the persistent config options (--nonpublished/--ncbi_only) are barred + # after initialization above, so a non-init run leaves the on-disk + # config untouched; `mtdb configure` is what edits it now + if init: if format_path(init) != format_path(os.environ["MYCODB"] + "../../"): - eprint("\nERROR: MTDB linked. Unlink via `mtdb -u`") + logger.error("MTDB linked. Unlink via `mtdb -u`") sys.exit(175) + elif args.init: + config = gen_config( + branch=kingdom, + forbidden="$MYCODB/log/forbidden.tsv", + nonpublished=nonpublished, + jgi=jgi, + rank2lineages=rank2lineages, + ) + write_json(config, init_dir + "config/mtdb.json", indent=1) - # nonfungi is nonpublished by default because it is all GenBank - if kingdom != "fungi": - nonpublished = True - # archaic placeholder for reference / rogue DB setup - elif nonpublished and rogue_bool: - nonpublished = validate_t_and_c(config) - else: - nonpublished = False + nonpublished = parse_nonpublished(kingdom, nonpublished, config) + config["nonpublished"] = nonpublished + return config - # branch = 'stable' - db_path = primaryDB() - if not resume or add: - date = datetime.now().strftime("%Y%m%d") - else: - date = str(resume) - if not ncbi_email: - ncbi_email, ncbi_api, jgi_email, jgi_pwd = loginCheck() - Entrez.email = ncbi_email + +def gather_login(ncbi_only): + + ncbi_api, jgi_email, jgi_pwd = login_check() if ncbi_api: Entrez.api_key = ncbi_api - if init: - dbtype = kingdom - init_dir = format_path(init) - if os.path.isdir(init_dir): - init_dir += "mycotoolsdb/" - if not init_dir.endswith("/"): - init_dir += "/" - envs = { - "MYCOFNA": init_dir + "data/fna", - "MYCOFAA": init_dir + "data/faa", - "MYCOGFF3": init_dir + "data/gff3", - "MYCODB": init_dir + "mtdb/", - } - os.environ["MYCODB"] = init_dir + "mtdb/" - output, config = initDB( - init_dir, - dbtype, - envs, - dbtype, - date=date, - rogue=rogue_bool, - nonpublished=nonpublished, - jgi=jgi, - repo=format_path(reference), - rank2lineages=rank2lineages, - ) - for env in envs: - os.environ[env] = envs[env] - orig_db = db2df(mtdb()) # initialize a new database - update_path = output + "log/" + date + "/" - if not os.path.isdir(update_path): - os.mkdir(update_path) - mtdb_initialize( - init_dir, init=True - ) # init_dir + 'config/mtdb.json', init = True) - else: - try: - output = format_path("$MYCODB/..") - except KeyError: - eprint("\nERROR: MTDB not linked. Link via `mtdb -i `", flush=True) - sys.exit(50) - update_path = output + "log/" + date + "/" - if not os.path.isdir(update_path): - os.mkdir(update_path) - if not True: # config['rogue']: # NEED TO MAKE THIS wget a particular URL - old_db = db2df(db_path) - shutil.move(db_path, update_path + os.path.basename(db_path)) - git_pull = subprocess.call( - [ - "git", - "pull", - "-C", - output + "mtdb", - config["repository"], - "-B", - branch, - ], - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - ) - new_db = db2df(primaryDB()) - orig_db = pd.concat([old_db, new_db.loc[~new_db["ome"].isin(old_db.index)]]) - else: - orig_db = db2df(primaryDB()) + if ncbi_only: + jgi_email = None + + return ncbi_api, jgi_email, jgi_pwd - orig_db = orig_db.dropna(subset=["ome"]) +def set_kingdom_options(config, jgi_email): if config["branch"] in {"prokaryote", "bacteria"}: - jgi = False + jgi_email = None group = "prokaryotes" - king = "bacteria" # NEED to make DB tools pull from this + taxon = "bacteria" # NEED to make DB tools pull from this rank = "superkingdom" elif config["branch"] in {"plants"}: - jgi = False + jgi_email = None group = "eukaryotes" - king = "viridiplantae" + taxon = "viridiplantae" rank = "kingdom" - # elif config['branch'] in {'protists'}: - # jgi = False - # group = 'eukaryotes' - # king = 'protists' - # rank = 'kingdom' elif config["branch"] in {"animals"}: - jgi = False + jgi_email = None group = "eukaryotes" - king = "metazoa" + taxon = "metazoa" rank = "kingdom" elif config["branch"] in {"archaea"}: - jgi = False + jgi_email = None group = "prokaryotes" - king = "Archaea" + taxon = "Archaea" rank = "superkingdom" else: - jgi = not ncbi_only group = "eukaryotes" - king = "fungi" + taxon = "fungi" rank = "kingdom" + return jgi_email, group, taxon, rank - if add or predb: # add predb2mtdb 2 master database - if predb: - add_predb = read_predb(predb_path) - addDB, init_failed = predb2mtdb( - add_predb, orig_db, update_path, cpus=cpu, remove=False, spacer="\t\t" - ) - if init_failed: - if not failed: - eprint("\nERROR: some genomes failed curation", flush=True) - sys.exit(23) - else: - eprint("\nWARNING: some genomes failed curation", flush=True) +def prep_predb_opts(predb, rerun_failed): + add_predb = read_predb(predb) + add_mtdb, init_failed = predb2mtdb( + add_predb, orig_db, update_path, cpus=cpu, remove=False, spacer="\t\t" + ) + if init_failed: + if not rerun_failed: + raise ValueError("some genomes failed curation") else: - addDB = mtdb(format_path(add)) - # we need full Paths for an addDB - gff_fail, fna_fail, faa_fail = False, False, False - if not all(os.path.isfile(format_path(x)) for x in addDB.reset_index()["gff3"]): - eprint("\nERROR: some GFF paths do not exist", flush=True) - gff_fail = [ - x - for x in addDB.reset_index()["gff3"] - if not os.path.isfile(format_path(x)) - ] - print(",".join(gff_fail), flush=True) - if not all(os.path.isfile(format_path(x)) for x in addDB.reset_index()["fna"]): - eprint("\nERROR: some FNA paths do not exist", flush=True) - fna_fail = [ - x - for x in addDB.reset_index()["fna"] - if not os.path.isfile(format_path(x)) - ] - print(",".join(fna_fail), flush=True) - if not all(os.path.isfile(format_path(x)) for x in addDB.reset_index()["faa"]): - eprint("\nERROR: some FAA paths do not exist", flush=True) - faa_fail = [ - x - for x in addDB.reset_index()["faa"] - if not os.path.isfile(format_path(x)) - ] - print(",".join(faa_fail), flush=True) - if gff_fail or fna_fail or faa_fail: - sys.exit(124) - - addDB["aquisition_date"] = [date for x in addDB["ome"]] - # make date the acquisition time - orig_mtdb = mtdb(primaryDB()) - update_path = format_path("$MYCODB/../" + "log/" + date + "/") - if not os.path.isdir(update_path): - os.mkdir(update_path) - shutil.copy(primaryDB(), update_path) - - tax_path = f"{update_path}../taxonomy.tsv" - tax_dicts = read_prev_tax(tax_path) - tax_dicts = gather_taxonomy( - addDB, - api_key=ncbi_api, - king=king, - rank=rank, - tax_dicts=tax_dicts, - output_path=tax_path, - ) - addDB, genus_dicts = assimilate_tax(addDB, tax_dicts) - addDB = check_add_mtdb(orig_mtdb, addDB, update_path, overwrite) - - write_forbid_omes(set(addDB["ome"]), format_path("$MYCODB/../log/relics.txt")) - - new_mtdb, update_omes = db2primary(addDB, orig_mtdb, save=True) - new_db_path = format_path("$MYCODB/" + date + ".mtdb") - - new_mtdb.df2db(new_db_path) - - if new_db_path != db_path: - if db_path: - os.remove(db_path) - return new_db_path + logger.warning("some genomes failed curation") + return add_mtdb + + +def add2mtdb(add_mtdb, date, ncbi_api, taxon, rank): + # we need full Paths for an add_mtdb + gff_fail, fna_fail, faa_fail = False, False, False + if not all(Path(format_path(x)).is_file() for x in add_mtdb.reset_index()["gff3"]): + logger.error("some GFF paths do not exist") + gff_fail = [ + x + for x in add_mtdb.reset_index()["gff3"] + if not Path(format_path(x)).is_file() + ] + logger.error(",".join(gff_fail)) + if not all(Path(format_path(x)).is_file() for x in add_mtdb.reset_index()["fna"]): + logger.error("some FNA paths do not exist") + fna_fail = [ + x + for x in add_mtdb.reset_index()["fna"] + if not Path(format_path(x)).is_file() + ] + logger.error(",".join(fna_fail)) + if not all(Path(format_path(x)).is_file() for x in add_mtdb.reset_index()["faa"]): + logger.error("some FAA paths do not exist") + faa_fail = [ + x + for x in add_mtdb.reset_index()["faa"] + if not Path(format_path(x)).is_file() + ] + logger.error(",".join(faa_fail)) + if gff_fail or fna_fail or faa_fail: + raise FileNotFoundError() + + add_mtdb["aquisition_date"] = [date for x in add_mtdb["ome"]] + # make date the acquisition time + orig_mtdb = mtdb(primary_db()) + update_path = format_path("$MYCODB/../" + "log/" + date + "/") + if not Path(update_path).is_dir(): + Path(update_path).mkdir() + shutil.copy(primary_db(), update_path) - if taxonomy: - new_db, update_mtdb = taxonomy_update( - orig_db, - update_path, - date, - config, - ncbi_email, - ncbi_api, - rank=rank, - group=king, - ) - new_path = format_path("$MYCODB/" + date + ".mtdb") - update_mtdb.df2db(new_path) - sys.exit(0) - elif reference: - if any(not x for x in ref_db["published"]) and not nonpublished: - eprint( - "\nWARNING: nonpublished data detected in reference and will be ignored", - flush=True, - ) + tax_path = f"{update_path}../taxonomy.tsv" + tax_dicts = read_prev_tax(tax_path) + tax_dicts = gather_taxonomy( + add_mtdb, + api_key=ncbi_api, + king=taxon, + rank=rank, + tax_dicts=tax_dicts, + output_path=tax_path, + ) + genus_dicts = add_mtdb.assimilate_tax(tax_dicts) + add_mtdb = check_add_mtdb(orig_mtdb, add_mtdb, update_path) - new_mtdb, update_mtdb = ref_update( - ref_db, - update_path, - date, - failed, - jgi_email, - jgi_pwd, - config, - ncbi_email, - ncbi_api, - cpus=cpu, - check_MD5=not bool(no_md5), - jgi=jgi, - group=group, - kingdom=king, - remove=not save, - taxonomy=True, - ncbi_fallback=fallback, - ) - else: - new_mtdb, update_mtdb = rogue_update( - orig_db, - update_path, - date, - failed, - jgi_email, - jgi_pwd, - config, - ncbi_email, - ncbi_api, - cpus=cpu, - check_MD5=not bool(no_md5), - jgi=jgi, - group=group, - kingdom=king, - remove=not save, - lineage_constraints=config["lineage_constraints"], - ncbi_fallback=fallback, - ) + write_forbid_omes(set(add_mtdb["ome"]), format_path("$MYCODB/../log/relics.txt")) - if not update_mtdb: - eprint("\nNo new data acquired", flush=True) + new_mtdb, update_omes = db2primary(add_mtdb, orig_mtdb, save=True) + write_primary(new_mtdb, date, update_path) - if not save: # add the predb2mtdb and remove files - # df2db(db, format_path('$MYCODB/' + date + '.mtdb')) - # output new database and new list of omes - eprint("\nMoving data into database", flush=True) - write_forbid_omes( - set(new_mtdb["ome"]), format_path("$MYCODB/../log/relics.txt") - ) +def write_update_mtdb(new_mtdb, update_mtdb, date, update_path): + # output new database and new list of omes + logger.info("Moving data into database") + write_forbid_omes( + set(new_mtdb["ome"]), format_path("$MYCODB/../log/relics.txt") + ) - new_path = format_path("$MYCODB/" + date + ".mtdb") - if format_path(db_path) == new_path: - shutil.copy(db_path, db_path + ".tmp") - full_mtdb, update_omes = db2primary( - update_mtdb, new_mtdb, save=False, combined=True - ) - full_mtdb.df2db(new_path + ".tmp") - try: - shutil.move(primaryDB(), update_path + os.path.basename(primaryDB())) - # move master database to log if it exists - except FileNotFoundError: - pass - shutil.move(new_path + ".tmp", new_path) - rm_raw_data(update_path) - eprint("\nMTDB update complete", flush=True) - # gen_algn_db( - # update_path, set(full_mtdb['ome']) - # ) - else: - # NEED to: insert note aboutrunning updatedb on predb - new_mtdb.df2db(format_path(update_path + date + ".mtdb")) - eprint( - f"\nUpdate ready for `mtdb u -a` at " - + f'{format_path(update_path + date + ".mtdb")}' - ) - # output new database and new list of omes + full_mtdb, update_omes = db2primary( + update_mtdb, new_mtdb, save=False, combined=True + ) + write_primary(full_mtdb, date, update_path) + rm_raw_data(update_path) + logger.info("MTDB update complete") - return primaryDB() def main(): @@ -2133,15 +1986,13 @@ def main(): help="[-u] Do not integrate/delete new data; -a to complete", ) - # init_args.add_argument('--reinit', action = 'store_true', help = 'Redownload all web data') - # parser.add_argument('--rogue', action = 'store_true', - # help = 'De novo MTDB') # currently required - - conf_args = parser.add_argument_group("Configuration") + conf_args = parser.add_argument_group( + "Configuration (initialization only; change later via `mtdb configure`)" + ) conf_args.add_argument( "--nonpublished", action="store_true", - help="[FUNGI]: Include MycoCosm restricted-use", + help="[FUNGI, -i]: Include MycoCosm restricted-use", ) conf_args.add_argument( "--ncbi_only", help="[FUNGI, -i]: Forego MycoCosm", action="store_true" @@ -2157,9 +2008,6 @@ def main(): conf_args.add_argument("--failed", action="store_true", help="Rerun/ignore failed") conf_args.add_argument("--forbidden", action="store_true", help="Rerun forbidden") - # conf_args.add_argument('--deviate', action = 'store_true', help = 'Deviate' \ - # + ' from existing config without prompting') - run_args = parser.add_argument_group("Runtime") run_args.add_argument("--resume", type=int, help="Resume previous date (YYYYmmdd)") run_args.add_argument( @@ -2168,57 +2016,176 @@ def main(): help="Skip NCBI MD5" + " (expedite large reruns)", ) run_args.add_argument( - "--fallback", - action="store_false", - default=True, - help="[ALPHA] use NCBI datasets utility for downloading NCBI data", + "--chunk", + type=int, + default=25, + help="Genomes to download per datasets call; larger chunks are " + + "likelier to be reset mid-transfer by NCBI; DEFAULT: 25", + ) + run_args.add_argument( + "--report_chunk", + type=int, + default=10000, + help="Metadata reports to acquire per datasets call; these are far " + + "smaller than genomes, so this is sized above --chunk; DEFAULT: 500", + ) + run_args.add_argument( + "--tape_wait", + type=int, + default=None, + help="[FUNGI]: Maximum minutes to wait for a MycoCosm genome's tape " + + "restore before deferring it to a later run; DEFAULT: wait indefinitely", ) run_args.add_argument("-c", "--cpu", type=int, default=1) + run_args.add_argument( + "-v", + "--verbose", + action="store_true", + help="Report per-genome JGI/NCBI search diagnostics (DEBUG logging)", + ) args = parser.parse_args() + setup_logging(verbose=getattr(args, "verbose", False)) args_dict = { - "Primary MTDB": primaryDB(verbose=False), + "Primary MTDB": primary_db(verbose=False), "Update": args.update, "Initialize": args.init, - "Add": format_path(args.add), #'Rogue': rogue_bool, + "Add": format_path(args.add), "Include Restricted": bool(args.nonpublished), "Resume": args.resume, "Retry failed": args.failed, "Retry forbidden": args.forbidden, "Save raw data": args.save, + "Genome chunk": args.chunk, + "Report chunk": args.report_chunk, + "Max tape wait": ( + "indefinite" if args.tape_wait is None else f"{args.tape_wait} minute(s)" + ), } - findExecs(["datasets"], exit={"datasets"}) + find_execs(["datasets"], exit={"datasets"}) start_time = intro("Update MycotoolsDB", args_dict) - control_flow( - args.init, - args.update, - args.reference, - args.add, - args.taxonomy, - args.predb, - args.save, - args.nonpublished, - args.ncbi_only, - args.lineage, - args.rank, - args.kingdom, - args.failed, - args.forbidden, - args.resume, - args.no_md5, - args.cpu, - ncbi_email=None, - overwrite=not args.keep, - fallback=args.fallback, - ) + error_handle_args(args) + kingdom = determine_kingdom(args.kingdom) + rank2lineages = parse_lineages(args.lineage, args.rank) + ncbi_api, jgi_email, jgi_pwd = gather_login(args.ncbi_only) + config = import_config(args.init, kingdom, args.nonpublished, bool(jgi_email), rank2lineages) - outro(start_time) + db_path = primary_db() + if not args.resume or args.add: + date = datetime.now().strftime("%Y%m%d") + else: + date = str(args.resume) + + # initialize + if args.init: + orig_db, update_path = init_db(format_path(args.init)) + # load existing database + else: + try: + output = format_path("$MYCODB/..") + except KeyError: + raise FileNotFoundError("MTDB not linked. Link via `mtdb -i `") + update_path = output + "log/" + date + "/" + orig_db = db2df(primary_db()) + + if not Path(update_path).is_dir(): + Path(update_path).mkdir() + orig_db = orig_db.dropna(subset=["ome"]) + jgi_email, group, taxon, rank = set_kingdom_options(config, jgi_email) + + # update taxonomy + if args.taxonomy: + update_mtdb = taxonomy_update( + orig_db, + update_path, + date, + config, + ncbi_api, + rank=rank, + group=taxon, + ) + write_primary(update_mtdb, date, update_path) + return 0 + + # add finalized DB + if args.add or args.predb: # add predb2mtdb 2 master database + if args.predb: + add_mtdb = prep_predb_opts(format_path(args.predb), args.failed) + else: + add_mtdb = mtdb(format_path(args.add)) + add2mtdb(add_mtdb, date, ncbi_api, taxon, rank) + return 0 + + # update w/a reference + if args.reference: + ref_db = mtdb(format_path(reference), add_paths=False) + if any(not x for x in ref_db["published"]) and not config["nonpublished"]: + logger.warning( + "nonpublished data detected in reference and will be ignored" + ) + new_mtdb, update_mtdb = ref_update( + ref_db, + update_path, + date, + failed, + jgi_email, + jgi_pwd, + config, + ncbi_api, + cpus=args.cpu, + check_MD5=not bool(args.no_md5), + group=group, + kingdom=taxon, + remove=not args.save, + taxonomy=True, + chunk=args.chunk, + tape_wait=args.tape_wait, + ) + # update de novo + else: + new_mtdb, update_mtdb = rogue_update( + orig_db, + update_path, + date, + args.failed, + jgi_email, + jgi_pwd, + config, + ncbi_api, + cpus=args.cpu, + check_MD5=not bool(args.no_md5), + group=group, + kingdom=taxon, + remove=not args.save, + lineage_constraints=config["lineage_constraints"], + chunk=args.chunk, + report_chunk=args.report_chunk, + tape_wait=args.tape_wait, + ) + + if not update_mtdb: + logger.info("No new data acquired") + return 0 + elif args.save: + new_mtdb.df2db(format_path(update_path + date + ".mtdb")) + logger.info( + f"Update ready for `mtdb u -a` at " + + f'{format_path(update_path + date + ".mtdb")}' + ) + return 0 + else: + write_update_mtdb(new_mtdb, update_mtdb, date, update_path) + return 0 def cli(): - main() + # BioPython (Bio.Entrez) raises a UserWarning when Entrez.email is unset; + # silence it so update output stays readable. + warnings.filterwarnings("ignore", category=UserWarning, module=r"Bio(\.|$)") + exit_code = main() + outro(start_time) if __name__ == "__main__": diff --git a/mycotools/ncbiAcc2fa.py b/mycotools/ncbiAcc2fa.py deleted file mode 100755 index 22ad287..0000000 --- a/mycotools/ncbiAcc2fa.py +++ /dev/null @@ -1,89 +0,0 @@ -#! /usr/bin/env python3 - -import os -import sys -import time -import getpass -from Bio import Entrez -from mycotools.lib.kontools import file2list, eprint, sys_start - - -def entrez_login(): - """Login to Entrez from user input""" - email = input("\nInput NCBI login email: ") - limit = 3 - Entrez.email = email - if len(accs) > 3: - api = getpass.getpass(prompt="NCBI API key (leave blank if none): ") - if api != "": - Entrez.api_key = api - limit = 10 - - eprint(flush=True) - return limit - - -def grab_accs(accs, limit): - """Grab FASTAs of NCBI accessions""" - count, amount, out_str = 0, 1, "" - for acc in accs: - count += 1 - # do not overwhelm the server - if count >= limit: - count = 0 - time.sleep(1) - eprint(acc, flush=True) - - # iteratively query until successful - attempt = 0 - while attempt < 3: - attempt += 1 - try: - handle = Entrez.efetch(db="protein", id=acc, retmode="xml") - records = Entrez.read(handle) - out_str += ( - f">{acc} " - + f"{records[0]['GBSeq_organism'].replace('','_')}\n" - + f"{records[0]['GBSeq_sequence'].upper()}\n" - ) - break - except: - time.sleep(1) - - return out_str - - -def cli(): - """Command line entrance""" - usage = ( - "Input NCBI accession or new line delimitted " - + "file of accessions and optionally the column name." - + "\nncbiAccs2fa.py \n" - ) - - # parse the arguments - args = sys_start(sys.argv, usage, 2) - - if len(args) <= 3: - # import a file of accessions - if os.path.isfile(args[1]): - if len(args) == 3: - accs = file2list(args[1], sep="\t", col=args[2]) - else: - accs = file2list(args[1]) - # import the command line accessions - else: - accs = [args[1]] - - limit = entrez_login() - out_str = grab_accs(accs, limit) - - eprint(flush=True) - with open(args + ".retr.fa", "w") as out: - out.write(out_str) - - sys.exit(0) - - -if __name__ == "__main__": - cli() diff --git a/mycotools/ncbi_dwnld_fallback.py b/mycotools/ncbi_dwnld_fallback.py deleted file mode 100644 index c774575..0000000 --- a/mycotools/ncbi_dwnld_fallback.py +++ /dev/null @@ -1,1004 +0,0 @@ -#! /usr/bin/env python3 - -# NEED a db check to ensure the log is relevant to the input -# NEED to convert to datasets -# NEED to consider refseq genomes with annotations when genbank doesn't have them - -import os -import re -import sys -import gzip -import time -import shutil -import urllib -import urllib.request -import requests -import argparse -import subprocess -import numpy as np -import pandas as pd -from contextlib import closing -from tqdm import tqdm -from Bio import Entrez -from datetime import datetime -from mycotools.lib.kontools import ( - intro, - outro, - format_path, - prep_output, - eprint, - vprint, - findExecs, -) -from mycotools.lib.dbtools import log_editor, loginCheck, mtdb, read_tax - - -def ncbidb2df(data, stdin=False): - import pandas as pd, pandas - - columns = mtdb.columns - if isinstance(data, mtdb): - db_df = pd.DataFrame(data.reset_index()) - elif not stdin: - data = format_path(data) - db_df = pd.read_csv(data, sep="\t") - if "ome" not in set(db_df.columns) and "assembly_acc" not in set(db_df.columns): - db_df = pd.read_csv(data, sep="\t", header=None) - else: - db_df = pd.read_csv(StringIO(data), sep="\t") - if "ome" not in set(db_df.columns) and "assembly_acc" not in set(db_df.columns): - db_df = pd.read_csv(StringIO(data), sep="\t", header=None) - - db_df = db_df.fillna("") - - return db_df - - -def prepare_folders(output_path, gff, prot, assem, transcript): - - file_types = [] - if assem: - if not os.path.exists(output_path + "fna"): - os.mkdir(output_path + "fna") - file_types.append("fna") - if gff: - if not os.path.exists(output_path + "gff3"): - os.mkdir(output_path + "gff3") - file_types.append("gff3") - if prot: - if not os.path.exists(output_path + "faa"): - os.mkdir(output_path + "faa") - file_types.append("faa") - if transcript: - if not os.path.exists(output_path + "transcript"): - os.mkdir(output_path + "transcript") - file_types.append("transcript") - - return file_types - - -def compile_log(output_path, remove=False): - - acc2log = {} - if not os.path.isfile(output_path + "ncbiDwnld.fallback.log"): - with open(output_path + "ncbiDwnld.fallback.log", "w") as out: - out.write( - "#ome\tassembly_acc\tassembly\tproteome\tgff3\ttranscript\t" - + "fna_md5\tfaa_md5\tgff3_md5\ttrans_md5\tgenome_id(s)\tgenus\tspecies\tstrain" - ) - - # too risky, too many things can go wrong and then users would be in a - # loop, but necessary for huge downloads - else: - with open(output_path + "ncbiDwnld.fallback.log", "r") as raw: - for line in raw: - if not line.startswith("#"): - data = [x.rstrip() for x in line.split("\t")] - while len(data) < 13: - data.append("") - acc2log[data[0]] = { - "assembly_acc": str(data[1]), - "fna": str(data[2]), - "faa": str(data[3]), - "gff3": str(data[4]), - "transcript": str(data[5]), - "fna_md5": str(data[6]), - "faa_md5": str(data[7]), - "gff3_md5": str(data[8]), - "transcript_md5": str(data[9]), - "genome_id": data[10], - "genus": str(data[11]), - "species": str(data[12]), - "strain": str(data[13]), - } - - return acc2log - - -def wait_for_ncbi(count, api=False): - if count >= 2: - if not api: - time.sleep(1) - count = 0 - elif count >= 7: - time.sleep(1) - count = 0 - return count - - -def esearch_ncbi(accession, column, database="assembly"): - search_term, esc_count = accession + "[" + column + "]", 0 - while esc_count < 3: - try: - handle = Entrez.esearch(db=database, term=search_term) - genome_ids = Entrez.read(handle)["IdList"] - break - except (RuntimeError, urllib.error.HTTPError) as e: - time.sleep(1) - esc_count += 1 - else: - print("\tERROR:", accession, "failed to search NCBI") - return None - return genome_ids - - -def esummary_ncbi(ID, database): - - esc_count = 0 - while esc_count < 10: - esc_count += 1 - try: - handle = Entrez.esummary(db=database, id=ID, report="full") - record = Entrez.read(handle, validate=False) - except urllib.error.HTTPError: - time.sleep(0.1) - continue - if database == "assembly": - try: # is it populated with an FTP? - ftp_path = str( - record["DocumentSummarySet"]["DocumentSummary"][0][ - "FtpPath_GenBank" - ] - ) - except IndexError: # wait a sec and retry - time.sleep(1) - continue - break - else: # too many failed attempts - if esc_count >= 10: - raise urllib.error.HTTPError("\tERROR: FTP request failed") - - return record - - -# collects paths to download proteomes and assemblies -def collect_ftps( - ncbi_df, - acc2log, - api_key=0, - column="assembly_acc", - ncbi_column="Assembly Accession", - database="assembly", - output_path="", - verbose=True, - remove=False, - spacer="\t\t", -): - - count, failed = 0, [] - - # for each row in the assembly, grab the accession number, form the search term for Entrez, use Entrez, - if ncbi_column in {"assembly", "genome", "uid"}: - out_df = ncbi_df[ncbi_df.index.isin(set(acc2log.keys()))] - ncbi_df = ncbi_df[~ncbi_df.index.isin(set(acc2log.keys()))] - for accession, row in tqdm(ncbi_df.iterrows(), total=len(ncbi_df)): - if accession in acc2log: # add all rows that have indices associated with this - # query type - out_df = pd.concat([out_df, row.to_frame().T]) - icount = 1 - test = str(accession) + "_" + str(icount) - while test in acc2log: - count += 1 - # row['assembly_acc'] = acc2log[test]['genome_id'] - if "ome" in row.keys(): - row["ome"] = None # haven't assigned a mycotools ID yet - out_df = pd.concat([out_df, row.to_frame().T]) - sys.exit() - test = str(accession) + "_" + str(icount) - continue - - elif pd.isnull(row[column]) or not row[column]: # ignore blank entries - acc2log[str(accession)] = { - "assembly_acc": accession, - "fna": "", - "faa": "", - "gff3": "", - "transcript": "", - "fna_md5": "", - "faa_md5": "", - "gff3_md5": "", - "transcript_md5": "", - "genome_id": "", - "genus": "", - "species": "", - "strain": "", - } - failed.append([accession, datetime.strftime(row["version"], "%Y%m%d")]) - continue - - if ncbi_column not in {"uid"}: # we already have the uid, no worries - genome_id = esearch_ncbi(accession, ncbi_column, database="assembly") - else: - genome_id = [accession] - - if not genome_id: # No IDs retrieved - if "ome" in row.keys(): - accession = row["ome"] - eprint(spacer + "\t" + accession + " failed to find genome ID", flush=True) - try: - failed.append([accession, datetime.strftime(row["version"], "%Y%m%d")]) - except TypeError: # if the row can't be formatted as a date - failed.append([accession, row["version"]]) - continue - - if ncbi_column in { - "Assembly Accession", - "assembly", - "genome", - "uid", - }: # be confident it is the most - # recent assembly UID - genome_id = [str(max([int(i) for i in genome_id]))] - - icount = 0 - for ID in genome_id: - if icount: - new_acc = str(accession) + "$" + str(icount) - else: - new_acc = accession - # obtain the path from a summary of the ftp directory and create the standard paths for proteomes and assemblies - acc2log[str(new_acc)] = { - "assembly_acc": accession, - "fna": "", - "faa": "", - "gff3": "", - "transcript": "", - "fna_md5": "", - "faa_md5": "", - "gff3_md5": "", - "transcript_md5": "", - "genome_id": ID, - "genus": "", - "species": "", - "strain": "", - } - record = esummary_ncbi(ID, database) - - strain = "" # populate a fallback strain - record_info = record["DocumentSummarySet"]["DocumentSummary"][0] - assemblyID = record_info["AssemblyAccession"] - org = record_info["SpeciesName"].split() - genus = org[0] - if len(org) > 2: - species = org[1] - strain1 = "".join(org[2:]) - elif len(org) == 2: - species = org[1] - strain1 = "" - else: - species = "sp." - strain1 = "" - - try: - for attr in record_info["Biosource"]["InfraspeciesList"]: - if attr["Sub_type"].lower() == "strain": - strain = attr["Sub_value"] - if strain.lower() in {"missing", "none"}: - strain = "" - break - except KeyError: - if strain1: - strain = strain1 - - strain = re.sub(r"[^a-zA-Z0-9]", "", strain) - - ftp_path = str(record_info["FtpPath_GenBank"]) - - if not ftp_path: - eprint( - spacer + "\t" + new_acc + " failed to return any FTP path", - flush=True, - ) - try: - failed.append( - [accession, datetime.strftime(row["version"], "%Y%m%d")] - ) - except TypeError: - failed.append([accession, str(row["version"])]) - continue - - esc_count = 0 - ass_md5, gff_md5, trans_md5, prot_md5, md5s = "", "", "", "", {} - basename = os.path.basename(ftp_path) - - dwnld = 0 - for attempt in range(3): - try: - r = requests.head( - ftp_path.replace("ftp://", "https://"), allow_redirects=True - ) - break - except: - time.sleep(1) - - if r.status_code != 200: - dwnld = -1 - else: - md5_path = ftp_path.replace("ftp://", "https://") + "/md5checksums.txt" - - dwnld = subprocess.call( - [ - "curl", - md5_path, - "-o", - output_path + ".tmpmd5", - "--connect-timeout", - "5", - ], - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - ) - - count += 1 - if dwnld == 0: - with open(output_path + ".tmpmd5", "r") as raw: - for line in raw: - data = line.rstrip().split(" ") - # data = line.rstrip().split() - if data and len(data) == 2: - try: - md5s[ftp_path + "/" + os.path.basename(data[1])] = data[ - 0 - ] - except IndexError: # 404 error or something else - md5s = {} - break - else: - md5s = {} - - tranname = os.path.basename(ftp_path.replace("/GCA", "/GCF")) - # tranname = os.path.basename(ftp_path) - assembly = ftp_path + "/" + basename + "_genomic.fna.gz" - if assembly in md5s: - ass_md5 = md5s[assembly] - else: - assembly = "" - proteome = ftp_path + "/" + basename + "_protein.faa.gz" - if proteome in md5s: - prot_md5 = md5s[proteome] - else: - proteome = "" - gff3 = ftp_path + "/" + basename + "_genomic.gff3.gz" - test_gff3 = re.sub(r"\.gff3\.gz$", ".gff.gz", gff3) - if gff3 in md5s: - gff_md5 = md5s[gff3] - elif test_gff3 in md5s: - gff3 = test_gff3 - gff_md5 = md5s[gff3] - else: - gff3 = "" - - transcript = ( - ftp_path.replace("/GCA", "/GCF") + "/" + tranname + "_rna.fna.gz" - ) - # transcript = f'{ftp_path}/{tranname}_rna.fna.gz' - if transcript in md5s: - trans_md5 = md5s[transcript] - - if (not assembly or not gff3) and remove: - try: - failed.append( - [accession, datetime.strftime(row["version"], "%Y%m%d")] - ) - except TypeError: - failed.append( - [accession, datetime.strftime(datetime.now(), "%Y%m%d")] - ) - - log_editor( - output_path + "ncbiDwnld.fallback.log", - str(new_acc), - str(accession) - + "\t" - + accession - + "\t" - + assembly - + "\t" - + proteome - + "\t" - + gff3 - + "\t" - + transcript - + "\t" - + ass_md5 - + "\t" - + prot_md5 - + "\t" - + gff_md5 - + "\t" - + trans_md5 - + "\t" - + ID - + f"\t{genus}\t{species}\t{strain}", - ) - acc2log[str(new_acc)] = { - "assembly_acc": accession, - "fna": assembly, - "fna_md5": ass_md5, - "faa": proteome, - "faa_md5": prot_md5, - "gff3": gff3, - "gff3_md5": gff_md5, - "transcript": transcript, - "transcript_md5": trans_md5, - "genome_id": ID, - "genus": genus, - "species": species, - "strain": strain, - } - row["dwnld_id"] = ID - if icount: - if "ome" in row: - row["ome"] = None - out_df = pd.concat([out_df, row.to_frame().T]) - icount += 1 - - # if no API key is used, we can only generate 3 queries per second, otherwise we can use 10 - count = wait_for_ncbi(count, api_key) - - return acc2log, failed, out_df - - -# download the file depending on the type inputted -def download_files( - acc_prots, acc, file_types, output_dir, count, remove=False, spacer="\t\t" -): - - dwnlds = {} - for file_type in file_types: - ftp_link = acc_prots[file_type] - dwnlds[file_type] = -1 - if file_type == "fna": - file_path = output_dir + "fna/" + os.path.basename(acc_prots[file_type]) - elif file_type == "gff3": - file_path = output_dir + "gff3/" + os.path.basename(acc_prots[file_type]) - elif file_type == "faa": - file_path = output_dir + "faa/" + os.path.basename(acc_prots[file_type]) - elif file_type == "transcript": - file_path = ( - output_dir + "transcript/" + os.path.basename(acc_prots[file_type]) - ) - - if os.path.isfile(file_path): - count += 1 - md5_cmd = subprocess.run( - ["md5sum", file_path], stdout=subprocess.PIPE - ) # check the md5 - md5_res = md5_cmd.stdout.decode("utf-8") - md5_find = re.search(r"\w+", md5_res) - md5 = md5_find[0] - - if md5 == acc_prots[file_type + "_md5"]: - eprint( - f"{spacer}\t{file_type}: {os.path.basename(file_path)}", flush=True - ) - dwnlds[file_type] = 0 - continue - elif os.path.isfile(file_path[:-3]): - eprint(f"{spacer}\t{file_type}: {os.path.basename(file_path)}", flush=True) - dwnlds[file_type] = 0 - continue - - if ftp_link == "": - dwnlds[file_type] = 15 - continue - # esc_count = 0 - for esc_count in range(3): - count += 1 - dwnld = subprocess.call( - ["curl", ftp_link, "-o", file_path + ".tmp", "--connect-timeout", "5"], - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - ) - if not dwnld: - os.rename(file_path + ".tmp", file_path) - break - else: - time.sleep(1) - count = 0 - - if dwnld: - eprint(f"{spacer}\t\tERROR: {file_type} failed", flush=True) - dwnlds[file_type] = 69 - acc_prots[file_type] = "" - log_editor( - output_dir + "ncbiDwnld.fallback.log", - str(acc), - str(acc) - + "\t" - + str(acc_prots["assembly_acc"]) - + "\t" - + str(acc_prots["fna"]) - + "\t" - + str(acc_prots["faa"]) - + "\t" - + str(acc_prots["gff3"]) - + "\t" - + str(acc_prots["transcript"]) - + "\t" - + str(acc_prots["fna_md5"]) - + "\t" - + str(acc_prots["faa_md5"]) - + "\t" - + str(acc_prots["gff3_md5"]) - + "\t" - + str(acc_prots["transcript_md5"]) - + "\t" - + f"{acc_prots['genome_id']}\t{acc_prots['genus']}\t" - + f"{acc_prots['species']}\t{acc_prots['strain']}", - ) - if remove and file_type in {"fna", "gff3"}: - break - continue - - if not os.path.isfile(file_path): - dwnlds[file_type] = 1 - eprint(f"{spacer}\t\tERROR: {file_type} missing", flush=True) - if remove and file_type in {"fna", "gff3"}: - break - else: - dwnlds[file_type] = 0 - if os.stat(file_path).st_size < 150: - eprint(f"{spacer}\t{file_type}: ERROR, file too small", flush=True) - dwnlds[file_type] = 420 - if remove and file_type in {"fna", "gff3"}: - break - eprint(f"{spacer}\t{file_type}: {os.path.basename(file_path)}", flush=True) - - return dwnlds, count - - -def dwnld_mngr(ncbi_df, data, acc, file_types, output_path, count, remove, api, spacer): - fail = [] - exits, count = download_files( - data, acc, file_types, output_path, count, remove=remove, spacer=spacer - ) - count = wait_for_ncbi(count, api) - t_acc = acc - try: - if exits["fna"] != 0: - if "$" in acc: - t_acc = acc[: acc.find("$")] - fail = [t_acc, ncbi_df["version"][t_acc]] - return fail, count - except KeyError: - pass - try: - if exits["gff3"] != 0: - if "$" in acc: # is this correct? - t_acc = acc[: acc.find("$")] - fail = [t_acc, ncbi_df["version"][t_acc]] - return fail, count - except KeyError: - pass - return fail, count - - -def dwnld_mngr_no_MD5( - ncbi_df, data, acc, file_types, output_path, count, remove, api, spacer -): - run, fail = False, [] - for file_type in file_types: - file_path = output_path + file_type + "/" + os.path.basename(data[file_type]) - if not os.path.isfile(file_path): - run = True - break - - if run: - exits, count = download_files( - data, acc, file_types, output_path, count, remove=remove, spacer=spacer - ) - count = wait_for_ncbi(count, api) - else: - exits = {x: 0 for x in file_types} - # we aren't checking md5s, so assume the exit is 0 - return fail, count - try: - if exits["fna"] != 0: - if "$" in acc: - t_acc = acc[: acc.find("$")] - else: - t_acc = acc - fail = [t_acc, ncbi_df["version"][t_acc]] - return fail, count - except KeyError: - pass - try: - if exits["gff3"] != 0: - if "$" in acc: # dont know if this should be here - t_acc = acc[: acc.find("$")] - else: - t_acc = acc - - fail = [t_acc, ncbi_df["version"][t_acc]] - return fail, count - except KeyError: - pass - return fail, count - - -def main( - api=None, - assembly=True, - proteome=False, - gff3=True, - transcript=False, - ncbi_df=False, - remove=False, - output_path=os.getcwd(), - verbose=False, - column="assembly_acc", - ncbi_column="Assembly", - check_MD5=True, - spacer="\t\t", -): - - # initialize run directory and information - output_path = format_path(output_path) - file_types = prepare_folders(output_path, gff3, proteome, assembly, transcript) - acc2log = compile_log(output_path, remove) - - # check if ncbi_df is a dataframe, and import if not - if not isinstance(ncbi_df, pd.DataFrame) and os.path.isfile(ncbi_df): - ncbi_df = ncbidb2df(ncbi_df) - if len(ncbi_df.index) == 0: - ncbi_df = pd.DataFrame({i: [v] for i, v in enumerate(list(ncbi_df.keys()))}) - - # make the modify date from a standard NCBI table the version if it does - # not otherwise exist, else there isn't a version to reference - if "Modify Date" in ncbi_df.keys() and not "version" in ncbi_df.keys(): - ncbi_df["version"] = pd.to_datetime(ncbi_df["Modify Date"]) - elif "version" not in ncbi_df.keys(): - ncbi_df["version"] = "" - - ncbi_df = ncbi_df.set_index(pd.Index(list(ncbi_df[column]))) - # preserve the original column, but index ncbi_df on it as well - vprint("\n" + spacer + "Assembling NCBI ftp directories", v=verbose, flush=True) - acc2log, failed, ncbi_df = collect_ftps( - ncbi_df, - acc2log, - remove=remove, - ncbi_column=ncbi_column, - column=column, - api_key=api, - output_path=output_path, - verbose=verbose, - spacer=spacer, - ) - - if remove: - acc2log = { - o: acc2log[o] - for o in acc2log - if all(acc2log[o][p] for p in ["fna", "gff3"]) - } - new_df = pd.DataFrame() - - vprint(f"\n{spacer}Downloading {len(acc2log)} NCBI files", v=verbose, flush=True) - count = 0 - if "strain" not in ncbi_df.columns: - ncbi_df["strain"] = "" - if check_MD5: - for acc, data in acc2log.items(): - eprint(spacer + "\t" + str(acc), flush=True) - if data: - fail, count = dwnld_mngr( - ncbi_df, - data, - acc, - file_types, - output_path, - count, - remove, - api, - spacer, - ) - if fail: - failed.append(fail) - else: - ncbi_df.at[acc, "assemblyPath"] = ( - output_path + "fna/" + os.path.basename(acc2log[acc]["fna"]) - ) - ncbi_df.at[acc, "faa"] = ( - output_path + "faa/" + os.path.basename(acc2log[acc]["faa"]) - ) - ncbi_df.at[acc, "gffPath"] = ( - output_path + "gff3/" + os.path.basename(acc2log[acc]["gff3"]) - ) - ncbi_df.at[acc, "genus"] = acc2log[acc]["genus"] - ncbi_df.at[acc, "species"] = acc2log[acc]["species"] - if not ncbi_df.loc[acc, "strain"] or pd.isnull( - ncbi_df.loc[acc, "strain"] - ): - ncbi_df.at[acc, "strain"] = acc2log[acc]["strain"] - new_df = pd.concat([new_df, ncbi_df.loc[acc].to_frame().T]) - else: - check = ncbi_df[ncbi_df[column] == acc[: acc.find("$")]] - # check for entries in the inputted table that match the accession - # provided without version modification - db_vers = datetime.strftime(row["version"], "%Y%m%d") - failed.append([acc, db_vers]) - else: # there is no checking md5, this is for efficient, so make it - # efficient by avoiding conditional expressions - for acc, data in acc2log.items(): - eprint(spacer + "\t" + str(acc), flush=True) - fail, count = dwnld_mngr_no_MD5( - ncbi_df, data, acc, file_types, output_path, count, remove, api, spacer - ) - if fail: - failed.append(fail) - else: - for file_type in file_types: - ncbi_df.at[acc, file_type] = ( - output_path - + file_type - + "/" - + os.path.basename(data[file_type]) - ) - new_df = pd.concat([new_df, ncbi_df.loc[acc].to_frame().T]) - - if "fna" in new_df.keys(): - new_df = new_df.rename(columns={"fna": "assemblyPath"}) - if "gff3" in new_df.keys(): - new_df = new_df.rename(columns={"gff3": "gffPath"}) - new_df = new_df.reset_index() - return new_df, failed - - -def get_SRA(assembly_acc, fastqdump="fastq-dump", pe=True): - - handle = Entrez.esearch(db="SRA", term=assembly_acc) - ids = Entrez.read(handle)["IdList"] - for id in ids: - handle = Entrez.esummary(db="SRA", id=id, report="full") - records = Entrez.read(handle, validate=False) - for record in records: - srr = re.search(r'Run acc="(S\w+\d+)"', record["Runs"])[1] - print("\t\t" + srr, flush=True) - cmd, count = 1, 0 - if pe: - while cmd and count < 3: - count += 1 - cmd = subprocess.call( - ["prefetch", srr, "--max-size", "10t"], stdout=subprocess.PIPE - ) - if cmd: - continue - cmd = subprocess.call(["vdb-validate", srr], stdout=subprocess.PIPE) - if cmd: - continue - cmd = subprocess.call( - [fastqdump, "--split-3", "--gzip", srr], stdout=subprocess.PIPE - ) - if os.path.isfile(srr + "_1.fastq"): - # if os.path.isfile(srr + '_1.fastq'): - # cmd = subprocess.call(['gzip', f'{srr}_1.fastq']) - # cmd = subprocess.call(['gzip', f'{srr}_2.fastq']) - shutil.move( - srr + "_1.fastq.gz", assembly_acc + "_" + srr + "_1.fq.gz" - ) - shutil.move( - srr + "_2.fastq.gz", assembly_acc + "_" + srr + "_2.fq.gz" - ) - else: - # cmd = subprocess.call(['gzip', f'{srr}.fastq']) - print( - "\t\t\tWARNING: file failed or not paired-end", flush=True - ) - else: - while cmd and count < 3: - count += 1 - cmd = subprocess.call(["prefetch", srr], stdout=subprocess.PIPE) - if cmd: - continue - cmd = subprocess.call(["vdb-validate", srr], stdout=subprocess.PIPE) - if cmd: - continue - cmd = subprocess.call( - [fastqdump, srr, "--gzip"], stdout=subprocess.PIPE - ) - if cmd: - continue - # cmd = subprocess.call(['gzip', f'{srr}.fastq']) - if os.path.isfile(srr + ".fastq.gz"): - shutil.move( - srr + ".fastq.gz", assembly_acc + "_" + srr + ".fq.gz" - ) - else: - print("\t\t\tERROR: file failed", flush=True) - - -def goSRA(df, output=os.getcwd() + "/", pe=True): - - print() - sra_dir = output + "sra/" - if not os.path.isdir(sra_dir): - os.mkdir(sra_dir) - os.chdir(sra_dir) - fastqdump = findExecs("fastq-dump", exit=set("fastq-dump")) - count = 0 - - if "sra" in df.keys(): - row_key = "sra" - else: - row_key = "assembly_acc" - - for i, row in df.iterrows(): - print("\t" + row[row_key], flush=True) - get_SRA(row[row_key], fastqdump[0]) - count += 1 - if count >= 10: - time.sleep(1) - count = 0 - - -def cli(): - parser = argparse.ArgumentParser( - description="GenBank downloading utility. Downloads " - + "accession by accession, files without MD5s are excluded" - ) - parser.add_argument( - "-i", - "--input", - required=True, - help="Space delimited accession; tab delimited file with -c", - ) - parser.add_argument("-a", "--assembly", action="store_true") - parser.add_argument("-p", "--proteome", action="store_true") - parser.add_argument("-g", "--gff3", action="store_true") - parser.add_argument("-t", "--transcript", action="store_true") - parser.add_argument("-s", "--sra", action="store_true", help="Download SRAs only") - parser.add_argument( - "-pe", - "--paired", - action="store_true", - help="Download paired-end SRAs. (REQUIRES -s)", - ) - parser.add_argument( - "-c", "--column", help='Accession column num/name; DEFAULT ["assembly_acc" | 0]' - ) - parser.add_argument( - "-n", - "--ncbi_column", - help="NCBI database associated with column. " - + '{"assembly", "biosample", "bioproject", "genome" ...}; ' - + "DEFAULT: attempt to decipher", - ) - parser.add_argument("-o", "--output", help="Output directory") - parser.add_argument("-e", "--email", help="NCBI email") - parser.add_argument("--api", help="NCBI API key for high query rate") - args = parser.parse_args() - - if args.email: - ncbi_email = args.email - Entrez.email = ncbi_email - if args.api: - ncbi_api = args.api - Entrez.api_key = ncbi_api - else: - ncbi_api = None - else: - ncbi_email, ncbi_api, jgi_email, jgi_pwd = loginCheck(jgi=False) - Entrez.email = ncbi_email - if ncbi_api: - Entrez.api_key = ncbi_api - - if not args.output: - output = os.getcwd() + "/" - else: - output = format_path(args.output) - - args_dict = { - "NCBI Table": args.input, - "email": ncbi_email, - "Assemblies": args.assembly, - "Proteomes": args.proteome, - ".gff3's": args.gff3, - "Transcripts": args.transcript, - "SRA": args.sra, - } - - start_time = intro("Download NCBI files", args_dict) - # if not args.assembly and not args.proteome and not args.gff3 and not args.sra and not args.transcript: - # eprint('\nERROR: You must choose at least one download option\nExit code 37', flush = True) - # sys.exit( 37 ) - - if args.sra: - if os.path.isfile(format_path(args.input)): - goSRA( - pd.read_csv(format_path(args.input), sep="\t"), output, pe=args.paired - ) - else: - goSRA(pd.DataFrame({"sra": [args.input.rstrip()]}), output, pe=args.paired) - else: - if os.path.isfile(format_path(args.input)): - ncbi_df = pd.read_csv(args.input, sep="\t", header=None) - if not args.column: - if "assembly_acc" in ncbi_df.keys(): - column = "assembly_acc" - ncbi_column = "Assembly Accession" - elif "Assembly Accession" in ncbi_df.keys(): - column = "assembly_acc" - ncbi_column = "Assembly Accession" - else: - column = 0 - else: - try: - column = ncbi_df.columns[int(0)] - ncbi_column = column - except ValueError: # not an integer - pass - if not args.ncbi_column: - if args.column is not None: - if column.lower() in {"assembly"}: - ncbi_column = "assembly" - elif column.lower() in { - "genome", - "assembly accession", - "assembly_acc", - }: - ncbi_column = "genome" - elif column.lower() in {"biosample", "biosample accession"}: - ncbi_column = "biosample" - else: - ncbi_column = column.lower() - else: - ncbi_column = "genome" - else: - ncbi_column = args.ncbi_column.lower() - else: - ncbi_df = pd.DataFrame( - {"assembly_acc": args.input.replace('"', "").replace("'", "").split()} - ) - column = "assembly_acc" - ncbi_column = "assembly" - - ncbi_df = ncbi_df.drop_duplicates(column) - - new_df, failed = main( - assembly=args.assembly, - column=column, - ncbi_column=ncbi_column, - proteome=args.proteome, - gff3=args.gff3, - transcript=args.transcript, - ncbi_df=ncbi_df, - output_path=output, - verbose=True, - spacer="", - ) - new_df = new_df.rename(columns={"index": "#assembly_accession"}) - new_df["source"] = "ncbi" - new_df["useRestriction (yes/no)"] = "no" - # if 'index' in new_df.columns: - # del new_df['index'] - if 0 in new_df.columns: - del new_df[0] - - new_df.to_csv(args.input + ".predb", sep="\t", index=None) - - outro(start_time) - - -if __name__ == "__main__": - cli() diff --git a/mycotools/phylo/__init__.py b/mycotools/phylo/__init__.py new file mode 100644 index 0000000..5c5069f --- /dev/null +++ b/mycotools/phylo/__init__.py @@ -0,0 +1,40 @@ +#! /usr/bin/env python3 +"""Dispatcher for the `mycotools phylo` subcommand. + +Routes `mycotools phylo ...` to a phylogenetics module.""" +from mycotools.lib.subcmd import Dispatcher + +# subcommand name/alias -> submodule within this package (mycotools.phylo.) +SUBCOMMANDS = { + "crap": "crap", + "tree": "tree", + "synteny": "synteny", + "tools": "tools", +} + +DESCRIPTION = """Build phylogenies and phylogenetic pipelines + +Tools (all following arguments are forwarded to the tool): + crap Cluster Reconstruction and Phylogeny (CRAP) pipeline + tree build a phylogeny from a fasta (align -> trim -> infer) + synteny build a microsynteny tree + tools manipulate an existing phylogeny (root/prune/rename/strip support) + +Examples: + mycotools phylo crap -h + mycotools phylo tools -h""" + +_dispatcher = Dispatcher( + "mycotools phylo", + "mycotools.phylo", + SUBCOMMANDS, + DESCRIPTION, + metavar="TOOL", + arg_help="phylogenetics tool (see below)", +) +main = _dispatcher.main +cli = _dispatcher.cli + + +if __name__ == "__main__": + cli() diff --git a/mycotools/phylo/__main__.py b/mycotools/phylo/__main__.py new file mode 100644 index 0000000..aca313d --- /dev/null +++ b/mycotools/phylo/__main__.py @@ -0,0 +1,6 @@ +#! /usr/bin/env python3 +"""Enable ``python -m mycotools.phylo`` to run the phylo dispatcher.""" +from mycotools.phylo import cli + +if __name__ == "__main__": + cli() diff --git a/mycotools/crap.py b/mycotools/phylo/crap.py similarity index 87% rename from mycotools/crap.py rename to mycotools/phylo/crap.py index 3ebb5b6..ec1b9ec 100755 --- a/mycotools/crap.py +++ b/mycotools/phylo/crap.py @@ -15,6 +15,7 @@ # NEED assembly reference method, i.e. tblastn # NEED cluster variable input default +import logging import os import re import sys @@ -28,42 +29,44 @@ from collections import Counter, defaultdict try: - from ete3 import Tree, faces, TreeStyle, NodeStyle, AttrFace + from ete3 import Tree, faces, TreeStyle, NodeStyle from ete3.parser.newick import NewickError except ImportError: raise ImportError( "Install ete3 into your conda environment via `conda install ete3`" ) -from mycotools.lib.dbtools import mtdb, primaryDB +from mycotools.lib.dbtools import mtdb, primary_db from mycotools.lib.kontools import ( - eprint, format_path, - findExecs, + find_execs, intro, outro, read_json, write_json, stdin2str, - getColors, + get_colors, collect_files, + setup_logging, ) -from mycotools.lib.biotools import fa2dict, dict2fa, gff2list, list2gff, gff3Comps -from mycotools.acc2fa import dbmain as acc2fa -from mycotools.fa2clus import ( - write_data, +from mycotools.lib.biotools import fa2dict, dict2fa, gff2list, list2gff, gff3_comps +from mycotools.mtdb.acc2.fa import dbmain as acc2fa +from mycotools.cluster.fasta import ( ClusteringError, ClusterParameterError, main as fa2clus, sort_iterations, ) -from mycotools.fa2tree import main as fa2tree, PhyloError -from mycotools.acc2locus import main as acc2locus -from mycotools.gff2svg import main as gff2svg -from mycotools.db2search import blast_main as db2search -from mycotools.ome2name import main as ome2name +from mycotools.phylo.tree import main as fa2tree, PhyloError +from mycotools.mtdb.acc2.locus import main as acc2locus +from mycotools.gff.svg import main as gff2svg +from mycotools.homology.db import blast_main as db2search +from mycotools.rename import main as ome2name # from mycotools.utils.og2mycodb import mycodbHGs, extract_ogs -from mycotools.db2microsyntree import compile_homolog_groups +from mycotools.phylo.synteny import compile_homolog_groups +from pathlib import Path + +logger = logging.getLogger(__name__) os.environ["QT_QPA_PLATFORM"] = "offscreen" @@ -172,8 +175,8 @@ def input_genes2input_hgs(input_genes, gene2hg): try: input_hgs[gene] = gene2hg[gene] except KeyError: - eprint("ERROR: " + gene + " query with no HG", flush=True) - eprint("\t" + gene + " will be ignored.", flush=True) + logger.error("" + gene + " query with no HG") + logger.info("" + gene + " will be ignored.") return input_hgs @@ -189,11 +192,11 @@ def check_fa_size(fas, max_size): submitted to tree building""" fas4clus, fas4trees = {}, {} for query, fa in fas.items(): - print("\t" + query + "\t" + str(len(fa)) + " genes", flush=True) + logger.debug("" + query + "\t" + str(len(fa)) + " genes") if len(fa) > max_size: fas4clus[query] = fa elif len(fa) < 3: - eprint(f"\t\tWARNING: too few hits ({len(fa)})", flush=True) + logger.warning(f"too few hits ({len(fa)})") else: fas4trees[query] = fa @@ -208,13 +211,13 @@ def write_seq_clus(gene_module, focal_gene, db, output_path, out_fa): try: fa_dict = acc2fa(db, gene_module) except KeyError: - eprint("\t\t\tWARNING: some hits not in database", flush=True) + logger.warning("some hits not in database") db_omes = set(db.keys()) gene_module = [x for x in gene_module if x[: x.find("_")] in db_omes] fa_dict = acc2fa(db, gene_module) # need to implement some method to choose if the max and min parameters couldn't be met - print("\t\t\t" + str(len(fa_dict)) + " genes in group", flush=True) + logger.debug("" + str(len(fa_dict)) + " genes in group") with open(out_fa, "w") as out: out.write(dict2fa(fa_dict)) @@ -247,15 +250,15 @@ def run_fa2clus( # output Diamond data to a diamond directory dmnd_dir = output + "dmnd/" - if not os.path.isdir(dmnd_dir): - os.mkdir(dmnd_dir) + if not Path(dmnd_dir).is_dir(): + Path(dmnd_dir).mkdir() # create an output path and a log path for the inputted gene output_path = output + str(focal_gene) log_path = output + "." + str(focal_gene) + ".log" # parse an existing fa2clus_log - if os.path.isfile(log_path): + if Path(log_path).is_file(): fa2clus_log = read_json(log_path) else: fa2clus_log = {"algorithm": "null"} @@ -282,9 +285,8 @@ def run_fa2clus( return False, False, fa2clus_log except ClusteringError as le: # if cluster error, try aggclus if not error: - eprint( - spacer + "mmseqs failed, attempting hierarchical " + "clustering", - flush=True, + logger.info( + spacer + "mmseqs failed, attempting hierarchical " + "clustering" ) try: cluster, newick, overshot, fa2clus_log = fa2clus( @@ -342,15 +344,12 @@ def outgroup_mngr( function to operate, or else an IndexError will result""" fa_path = clus_dir + focal_gene + ".fa" - out_name = str(focal_gene) + ".outgroup" output_path = clus_dir + str(focal_gene) dmnd_dir = clus_dir + "dmnd/" log_path = clus_dir + "." + str(focal_gene) + ".log" fa2clus_log = read_json(log_path) algorithm = fa2clus_log["algorithm"] # use previous search algorithm successes = fa2clus_log["successes"] - iterations = fa2clus_log["iterations"] - prev_size = len(fa2dict(clus_dir + "../" + str(focal_gene) + ".fa")) max_success = successes[0] # it is best to have the largest cluster, so attempt to refine upward @@ -378,7 +377,7 @@ def outgroup_mngr( min_var=min_var, max_var=max_var, ) # aggclus will have a higher minimum connectivity - except ClusterParameterError as e: # has run previously refined all it + except ClusterParameterError: # has run previously refined all it pass else: @@ -467,7 +466,7 @@ def outgroup_mngr( min_var=min_var, max_var=max_var, ) # mmseqs - except ClusterParameterError as e: + except ClusterParameterError: pass [ @@ -507,23 +506,23 @@ def make_output(base_dir, new_log): # eprint('\nERROR: base output directory missing: ' + base_dir, flush = True) # sys.exit(2) curdate = datetime.datetime.now().strftime("%Y%m%d") - output_dir = os.getcwd() + "/crap_" + curdate + "/" - if not os.path.isdir(output_dir): - os.mkdir(output_dir) + output_dir = str(Path.cwd()) + "/phylo_crap_" + curdate + "/" + if not Path(output_dir).is_dir(): + Path(output_dir).mkdir() else: - if not os.path.isdir(base_dir): - os.mkdir(base_dir) + if not Path(base_dir).is_dir(): + Path(base_dir).mkdir() output_dir = format_path(base_dir) # initialize a log in the output directory log_path = output_dir + ".craplog.json" parse_log(log_path, new_log, output_dir) - if not os.path.isdir(output_dir): - os.mkdir(output_dir) + if not Path(output_dir).is_dir(): + Path(output_dir).mkdir() loc_dir = output_dir + "loci/" - if not os.path.isdir(loc_dir): - os.mkdir(loc_dir) + if not Path(loc_dir).is_dir(): + Path(loc_dir).mkdir() wrk_dir = output_dir + "working/" global svg_dir # needs to be global for etetree @@ -532,12 +531,12 @@ def make_output(base_dir, new_log): svg_dir = wrk_dir + "svg/" gff_dir = wrk_dir + "genes/" tre_dir = wrk_dir + "trees/" - if not os.path.isdir(wrk_dir): - os.mkdir(wrk_dir) - if not os.path.isdir(svg_dir): - os.mkdir(svg_dir) - if not os.path.isdir(gff_dir): - os.mkdir(gff_dir) + if not Path(wrk_dir).is_dir(): + Path(wrk_dir).mkdir() + if not Path(svg_dir).is_dir(): + Path(svg_dir).mkdir() + if not Path(gff_dir).is_dir(): + Path(gff_dir).mkdir() return output_dir, wrk_dir, gff_dir, tre_dir @@ -593,20 +592,20 @@ def extract_locus_hg( try: gff_list = gff2list(gff3) except FileNotFoundError: - eprint(f"\t\t\tWARNING: {ome} MTDB entry without GFF3", flush=True) + logger.warning(f"{ome} MTDB entry without GFF3") return # grab the locus for the genes that do not currently have a locus SVG # (and thus also not a locus GFF) genes_to_grab = [ - x for x in genes_to_grab if not os.path.isfile(f"{wrk_dir}svg/{x}.locus.svg") + x for x in genes_to_grab if not Path(f"{wrk_dir}svg/{x}.locus.svg").is_file() ] try: out_indices, rna_gff = acc2locus( gff_list, genes_to_grab, plusminus, mycotools=True, geneGff=True, nt=True ) except ValueError: # KeyError: - eprint(f"\t\t\tWARNING: {ome} could not parse GFF", flush=True) + logger.warning(f"{ome} could not parse GFF") return # parse the loci and the genes associated with each locus, with the @@ -667,20 +666,20 @@ def extract_locus_gene( try: gff_list = gff2list(gff3) except FileNotFoundError: - eprint("\t\t\tWARNING: " + ome + " mycotoolsdb entry without GFF3", flush=True) + logger.warning("" + ome + " mycotoolsdb entry without GFF3") return - accs = [x for x in accs if not os.path.isfile(wrk_dir + "svg/" + x + ".locus.svg")] + accs = [x for x in accs if not Path(wrk_dir + "svg/" + x + ".locus.svg").is_file()] try: out_indices, rna_gff = acc2locus( gff_list, accs, plusminus, mycotools=True, geneGff=True, nt=True ) except KeyError: - eprint("\t\t\tWARNING: " + ome + " could not parse gff", flush=True) + logger.warning("" + ome + " could not parse gff") return extracted_genes, final_loci = {}, set() for locus_id, genes in out_indices.items(): - if os.path.isfile(wrk_dir + "svg/" + locus_id + ".locus.svg"): + if Path(wrk_dir + "svg/" + locus_id + ".locus.svg").is_file(): # NEED to rerun if there's a change in output parameters continue start_i, end_i = None, None @@ -757,7 +756,7 @@ def svgs2tree( circular=False, ): # svg_dir, out_dir): - init_dir = os.getcwd() + init_dir = str(Path.cwd()) os.chdir(out_dir) tree = Tree(tree_data) @@ -806,8 +805,8 @@ def svgs2tree( tree.render(out_dir + input_gene + "." + adj + ext, w=800, tree_style=ts) except TypeError: # QStandardPaths doesn't have permissions # this does not work - if not os.path.isdir(out_dir + ".XDG/"): - os.mkdir(out_dir + ".XDG/") + if not Path(out_dir + ".XDG/").is_dir(): + Path(out_dir + ".XDG/").mkdir() os.environ["XDG_RUNTIME_DIR"] = out_dir + ".XDG/" if og is not None: tree.render( @@ -830,16 +829,14 @@ def merge_color_palette(merges, query2color): def extend_color_palette(hgs, color_dict): hgs = [str(x) for x in hgs] new_hgs = set(hgs).difference(set(color_dict.keys())) - colors = getColors(len(hgs)) + colors = get_colors(len(hgs)) new_colors = list(set(colors).difference(set(color_dict.values()))) cor_i = 0 for i, v in enumerate(list(new_hgs)): try: color_dict[str(v)] = new_colors[i - cor_i] except IndexError: - eprint( - "\t\t\tWARNING: input too large for discrete arrow colors", flush=True - ) + logger.warning("input too large for discrete arrow colors") cor_i = i new_colors = colors color_dict[str(v)] = new_colors[i - cor_i] @@ -914,7 +911,7 @@ def make_color_palette(inputs, conversion_dict={}): for i, v in enumerate(inputs): color_dict[conversion_dict[v]] = extColors[i] else: - eprint("\nWARNING: input too large for discrete arrow colors", flush=True) + logger.warning("input too large for discrete arrow colors") try: for i, v in enumerate(inputs): color_dict[conversion_dict[v]] = extColors[i] @@ -938,7 +935,7 @@ def tree_mngr( ): query_fa_path = wrk_dir + query + ".fa" - if os.path.isfile(tre_dir + str(query) + tree_suffix) and reoutput: + if Path(tre_dir + str(query) + tree_suffix).is_file() and reoutput: return else: try: @@ -973,36 +970,31 @@ def parse_log(log_path, new_log, out_dir): except FileNotFoundError: old_log = None - rereun_search = False if old_log: try: - with open(old_log["db_path"], "rb") as raw: - db5 = hashlib.md5(raw.read()).hexdigest() - # if db5 != old_log['db']: if md5 changes do somethign - # rerun_search = True if old_log["search"] != new_log["search"]: - if os.path.isdir(out_dir): + if Path(out_dir).is_dir(): shutil.rmtree(out_dir) return elif old_log["bitscore"] != new_log["bitscore"]: fas = collect_files(wrk_dir, "fa") for fa in fas: - os.remove(fa) + Path(fa).unlink() fas = collect_files(clus_dir, "fa") for fa in fas: - os.remove(fa) - if os.path.isdir(wrk_dir + "tree/"): + Path(fa).unlink() + if Path(wrk_dir + "tree/").is_dir(): shutil.rmtree(wrk_dir + "tree/") elif old_log["plusminus"] != new_log["plusminus"]: - if os.path.isdir(wrk_dir + "genes/"): + if Path(wrk_dir + "genes/").is_dir(): shutil.rmtree(wrk_dir + "genes/") - if os.path.isdir(wrk_dir + "svg/"): + if Path(wrk_dir + "svg/").is_dir(): shutil.rmtree(wrk_dir + "svg/") elif old_log["labels"] != new_log["labels"]: - if os.path.isdir(wrk_dir + "svg/"): + if Path(wrk_dir + "svg/").is_dir(): shutil.rmtree(wrk_dir + "svg/") except KeyError: - eprint("\tERROR: log file corrupted. Hoping for the best.", flush=True) + logger.error("log file corrupted. Hoping for the best.") write_json(new_log, log_path) @@ -1039,7 +1031,7 @@ def crap_mngr( if info: return query2color - print("\t\tExtracting loci and generating synteny diagrams", flush=True) + logger.info("Extracting loci and generating synteny diagrams") extract_loci_cmds = [] if hg: for ome, ome_genes2hg in genes2query.items(): @@ -1089,9 +1081,9 @@ def crap_mngr( ] ) with mp.Pool(processes=cpus) as pool: - gene_res = pool.starmap(extract_locus_gene, extract_loci_cmds) + pool.starmap(extract_locus_gene, extract_loci_cmds) - print("\t\tMapping synteny diagrams on phylogeny", flush=True) + logger.info("Mapping synteny diagrams on phylogeny") tree_file = tre_dir + query + tree_suffix with open(tree_file, "r") as raw: raw_tree = raw.read() @@ -1117,7 +1109,7 @@ def crap_mngr( circular=circular, ) except NewickError: - eprint("\t\t\tERROR: newick malformatted", flush=True) + logger.error("newick malformatted") else: try: svgs2tree( @@ -1133,7 +1125,7 @@ def crap_mngr( circular=circular, # svg_dir, out_dir ) except NewickError: - eprint("\t\t\tERROR: newick malformatted", flush=True) + logger.error("newick malformatted") return query2color @@ -1146,7 +1138,7 @@ def write_loci(ome2genes, loc_dir): queries = ",".join(query) out.write(f"{gene}\t{queries}\n") else: - eprint(f"\t{ome} no hits", flush=True) + logger.info(f"{ome} no hits") # out.write('\n'.join(genes)) @@ -1206,7 +1198,7 @@ def parse_search_col_loci( id_mean = sum(identities) / len(identities) loc2sim[f] = [alia, locus_sim * id_mean] except FileNotFoundError: - eprint(f"\tWARNING: {ome} locus not weighted by %ID", flush=True) + logger.warning(f"{ome} locus not weighted by %ID") loc2sim[f] = [alia, locus_sim] # sort the loci by similarity @@ -1277,9 +1269,8 @@ def locus_output_mngr( for each genome via overlapping homology group similarity""" files = collect_files(gff_dir, "genes") ome2files = defaultdict(list) - ome2locs = defaultdict(list) for f in files: - ome = os.path.basename(f)[: os.path.basename(f).find("_")] + ome = Path(f).name[: Path(f).name.find("_")] ome2files[ome].append(f) if not report_dir: @@ -1349,17 +1340,14 @@ def hg_main( db = db.set_index() - print("\nCompiling homolog data", flush=True) - print("\tCompiling homologs", flush=True) + logger.info("Compiling homolog data") + logger.info("Compiling homologs") # og_info_dict = mycodbHGs(omes = set(db['ome'])) # hg2gene, gene2hg = extract_ogs(og_info_dict, ogtag) ome2i, gene2hg, i2ome, hg2gene = compile_homolog_groups( hg_file, wrk_dir, useableOmes=set(db.keys()) ) input_hgs = input_genes2input_hgs(input_genes, gene2hg) - input_hg2gene = { - v: k for k, v in input_hgs.items() - } # create hashes for transitioning todel, hits = [], set() for i, hg in enumerate(input_hgs): @@ -1375,12 +1363,12 @@ def hg_main( # in the future, genes without HGs will be placed into HGs via RBH if not input_hgs: - eprint("\nERROR: no HGs for any inputted genes", flush=True) + logger.error("no HGs for any inputted genes") sys.exit(3) if output_loci: - if all(os.path.isfile(f"{tre_dir}{gene}{tree_suffix}") for gene in input_hgs): - print("\nSkipping to outputting most similar loci to query", flush=True) + if all(Path(f"{tre_dir}{gene}{tree_suffix}").is_file() for gene in input_hgs): + logger.info("Skipping to outputting most similar loci to query") locus_output_mngr( gff_dir, loc_dir, @@ -1393,60 +1381,58 @@ def hg_main( return hg_fas = {} - if not all(os.path.isfile(f"{faa_dir}{hg}.faa") for gene, hg in input_hgs.items()): - print("\tPreparing homolog fastas", flush=True) + if not all(Path(f"{faa_dir}{hg}.faa").is_file() for gene, hg in input_hgs.items()): + logger.info("Preparing homolog fastas") compile_hg_fa_cmds = [ [db, hg2gene[hg], gene] for gene, hg in input_hgs.items() - if not os.path.isfile(wrk_dir + gene + ".fa") + if not Path(wrk_dir + gene + ".fa").is_file() ] with mp.Pool(processes=cpus) as pool: hg_fas = { x[0]: x[1] for x in pool.starmap(compile_hg_fa, compile_hg_fa_cmds) } for gene, hg in input_hgs.items(): - if os.path.isfile( + if Path( wrk_dir + gene + ".fa" - ): # add finished in working directory back + ).is_file(): # add finished in working directory back hg_fas = {**hg_fas, **{gene: fa2dict(wrk_dir + gene + ".fa")}} else: for gene, hg in input_hgs.items(): hg_fas = {**hg_fas, **{gene: fa2dict(f"{faa_dir}{hg}.faa")}} for gene in input_hgs: if gene not in hg_fas[gene]: - eprint( - f"\t\tERROR: {gene} not in homologs. Incorrect input?", flush=True - ) + logger.error(f"{gene} not in homologs. Incorrect input?") sys.exit(3) - print("\nChecking fasta sizes", flush=True) + logger.info("Checking fasta sizes") fas4clus, fas4trees = check_fa_size(hg_fas, max_size) for query, hit_fa in fas4trees.items(): hit_fa_path = wrk_dir + query + ".fa" with open(hit_fa_path, "w") as out: out.write(dict2fa(hit_fa)) - if not os.path.isdir(clus_dir): - os.mkdir(clus_dir) + if not Path(clus_dir).is_dir(): + Path(clus_dir).mkdir() if fas4clus: for query, fa in fas4clus.items(): clus_fa_path = clus_dir + query + ".fa" - if not os.path.isfile(clus_fa_path): + if not Path(clus_fa_path).is_file(): with open(clus_fa_path, "w") as out: out.write(dict2fa(fa)) - print("\nRunning clustering on " + str(len(fas4clus)) + " fastas", flush=True) + logger.debug("Running clustering on " + str(len(fas4clus)) + " fastas") - print("\nCRAP", flush=True) + logger.info("CRAP") ome_gene2hg = gene2hg2ome2hg(gene2hg) fas4trees = {k: v for k, v in sorted(fas4trees.items(), key=lambda x: len(x[1]))} for query, query_fa in fas4trees.items(): out_keys = None query_hits = list(query_fa.keys()) - print("\tQuery: " + str(query), flush=True) + logger.debug("Query: " + str(query)) if outgroups: - print("\t\tOutgroup detection", flush=True) - if os.path.isfile(clus_dir + query + ".fa"): - if not os.path.isfile(wrk_dir + query + ".outgroup.fa"): + logger.info("Outgroup detection") + if Path(clus_dir + query + ".fa").is_file(): + if not Path(wrk_dir + query + ".outgroup.fa").is_file(): in_keys, all_keys = outgroup_mngr( db, query, @@ -1459,14 +1445,13 @@ def hg_main( interval=interval, verbose=False, ) - out_query = query + ".outgroup" query_hits = all_keys out_keys = list(set(all_keys).difference(set(in_keys))) - print("\t\t\t" + str(len(in_keys)) + " gene ingroup", flush=True) + logger.debug("" + str(len(in_keys)) + " gene ingroup") if out_keys: - print("\t\t\t" + str(len(out_keys)) + " gene outgroup", flush=True) + logger.debug("" + str(len(out_keys)) + " gene outgroup") else: - eprint("\t\t\tWARNING: Could not detect outgroup for root", flush=True) + logger.warning("Could not detect outgroup for root") HG = input_hgs[ re.sub(r"\.outgroup$", "", query) ] # bulletproof against outgroups @@ -1496,8 +1481,8 @@ def hg_main( ) for query in fas4clus: - print("\tQuery: " + str(query), flush=True) - print("\t\tSequence clustering", flush=True) + logger.debug("Query: " + str(query)) + logger.info("Sequence clustering") out_keys = None res, overshot, fa2clus_log = run_fa2clus( clus_dir + query + ".fa", @@ -1515,11 +1500,11 @@ def hg_main( algorithm=clus_meth, ) if not res: - print("\t\t\tERROR: query had no significant hits", flush=True) + logger.error("query had no significant hits") continue if outgroups and not overshot: - print("\t\tOutgroup detection", flush=True) - if not os.path.isfile(wrk_dir + query + ".outgroup.fa"): + logger.info("Outgroup detection") + if not Path(wrk_dir + query + ".outgroup.fa").is_file(): in_keys, all_keys = outgroup_mngr( db, query, @@ -1532,12 +1517,11 @@ def hg_main( interval=interval, verbose=False, ) - out_query = query + ".outgroup" query_hits = all_keys out_keys = list(set(all_keys).difference(in_keys)) - print("\t\t\t" + str(len(in_keys)) + " gene ingroup", flush=True) + logger.debug("" + str(len(in_keys)) + " gene ingroup") if out_keys: - print("\t\t\t" + str(len(out_keys)) + " gene outgroup", flush=True) + logger.debug("" + str(len(out_keys)) + " gene outgroup") else: query_fa = fa2dict(wrk_dir + query + ".fa") query_hits = list(query_fa.keys()) @@ -1570,7 +1554,7 @@ def hg_main( ) if output_loci: - print("\nOutputting most similar loci to query", flush=True) + logger.info("Outputting most similar loci to query") locus_output_mngr( gff_dir, loc_dir, @@ -1616,7 +1600,7 @@ def search_main( ): """input_genes is a list of genes within an inputted cluster""" - print("\nPreparing run", flush=True) + logger.info("Preparing run") wrk_dir, loc_dir = out_dir + "working/", out_dir + "loci/" gff_dir, tre_dir = wrk_dir + "genes/", wrk_dir + "trees/" clus_dir = wrk_dir + "clus/" @@ -1628,23 +1612,23 @@ def search_main( query2color, conversion_dict = make_color_palette(input_genes, conversion_dict) if query_gff: - print("\tCleaning input GFF", flush=True) + logger.info("Cleaning input GFF") par_dict, prot_hits, RNA, query_gff = prep_gff( - query_gff, set(input_genes), gff3Comps() + query_gff, set(input_genes), gff3_comps() ) count = 0 while par_dict and count < 4: count += 1 par_dict, prot_hits, RNA, query_gff = prep_gff( - query_gff, set(input_genes), gff3Comps(), prot_hits, par_dict + query_gff, set(input_genes), gff3_comps(), prot_hits, par_dict ) if par_dict: - eprint("\nIncorrectly formatted GFF", flush=True) + logger.info("Incorrectly formatted GFF") sys.exit(5) elif set(input_genes).difference(prot_hits): - eprint("\nProteins missing from GFF", flush=True) - eprint( - "\t" + logger.info("Proteins missing from GFF") + logger.info( + "" + ",".join( [str(x) for x in list(set(input_genes).difference(prot_hits))] ) @@ -1663,22 +1647,22 @@ def search_main( query_fa = acc2fa(db, input_genes) with open(query_path, "w") as out: out.write(dict2fa(query_fa)) - elif not os.path.isfile(query_path): + elif not Path(query_path).is_file(): with open(query_path, "w") as out: out.write(dict2fa(query_fa)) search_fas = {} for query in query_fa: - if os.path.isfile(clus_dir + query + ".fa"): + if Path(clus_dir + query + ".fa").is_file(): search_fas[query] = fa2dict(clus_dir + query + ".fa") search_fas[query][query] = query_fa[query] - elif os.path.isfile(wrk_dir + query + ".fa"): + elif Path(wrk_dir + query + ".fa").is_file(): search_fas[query] = fa2dict(wrk_dir + query + ".fa") search_fas[query][query] = query_fa[query] if output_loci: - if all(os.path.isfile(f"{tre_dir}{gene}{tree_suffix}") for gene in query_fa): - print("\nSkipping to outputting most similar loci to query", flush=True) + if all(Path(f"{tre_dir}{gene}{tree_suffix}").is_file() for gene in query_fa): + logger.info("Skipping to outputting most similar loci to query") locus_output_mngr( gff_dir, loc_dir, @@ -1691,7 +1675,6 @@ def search_main( return skips = list(search_fas.keys()) - omes = set(db["ome"]) if not len(search_fas) == len(query_fa): if binary == "diamond": binary = "blastp" @@ -1720,14 +1703,14 @@ def search_main( for query in search_fas: search_fas[query][query] = query_fa[query] - print("\nChecking hit fasta sizes", flush=True) + logger.info("Checking hit fasta sizes") genes2query, merges = compile_genes_by_omes( search_fas, conversion_dict, set(db["ome"]) ) query2color = merge_color_palette(merges, query2color) for query in search_fas: # revert back to other fas - if os.path.isfile(wrk_dir + query + ".fa"): + if Path(wrk_dir + query + ".fa").is_file(): search_fas[query] = fa2dict(wrk_dir + query + ".fa") fas4clus, fas4trees = check_fa_size(search_fas, max_size) for query, hit_fa in fas4trees.items(): @@ -1736,26 +1719,25 @@ def search_main( out.write(dict2fa(hit_fa)) if fas4clus: - if not os.path.isdir(clus_dir): - os.mkdir(clus_dir) + if not Path(clus_dir).is_dir(): + Path(clus_dir).mkdir() for query, fa in fas4clus.items(): clus_fa_path = clus_dir + query + ".fa" - if not os.path.isfile(clus_fa_path): + if not Path(clus_fa_path).is_file(): with open(clus_fa_path, "w") as out: out.write(dict2fa(fa)) - print("\tRunning clustering on " + str(len(fas4clus)) + " fastas", flush=True) + logger.debug("Running clustering on " + str(len(fas4clus)) + " fastas") - print("\nCRAP", flush=True) - ome2genes = {} + logger.info("CRAP") fas4trees = {k: v for k, v in sorted(fas4trees.items(), key=lambda x: len(x[1]))} for query, query_fa in fas4trees.items(): out_keys = None query_hits = list(query_fa.keys()) - print("\tQuery: " + str(query), flush=True) + logger.debug("Query: " + str(query)) if outgroups: - print("\t\tOutgroup detection", flush=True) - if os.path.isfile(clus_dir + query + ".fa"): - if not os.path.isfile(wrk_dir + query + ".outgroup.fa"): + logger.info("Outgroup detection") + if Path(clus_dir + query + ".fa").is_file(): + if not Path(wrk_dir + query + ".outgroup.fa").is_file(): in_keys, all_keys = outgroup_mngr( db, query, @@ -1768,15 +1750,14 @@ def search_main( interval=interval, verbose=False, ) - out_query = query + ".outgroup" query_hits = all_keys out_keys = list(set(all_keys).difference(set(in_keys))) - print("\t\t\t" + str(len(in_keys)) + " gene ingroup", flush=True) + logger.debug("" + str(len(in_keys)) + " gene ingroup") if out_keys: - print("\t\t\t" + str(len(out_keys)) + " gene outgroup", flush=True) + logger.debug("" + str(len(out_keys)) + " gene outgroup") else: - eprint("\t\t\tWARNING: could not detect outgroup for root", flush=True) - null = crap_mngr( + logger.warning("could not detect outgroup for root") + crap_mngr( db, query, query_hits, @@ -1806,8 +1787,8 @@ def search_main( clus_cpus = cpus for query in fas4clus: - print("\tQuery: " + str(query), flush=True) - print("\t\tSequence clustering", flush=True) + logger.debug("Query: " + str(query)) + logger.info("Sequence clustering") out_keys = None query_hits = list(query_fa.keys()) @@ -1827,12 +1808,12 @@ def search_main( algorithm=clus_meth, ) if not res: - print("\t\t\tFAILED: query had no significant hits", flush=True) + logger.warning("FAILED: query had no significant hits") continue if outgroups and not overshot: - print("\t\tOutgroup detection", flush=True) - if not os.path.isfile(wrk_dir + query + ".outgroup.fa"): + logger.info("Outgroup detection") + if not Path(wrk_dir + query + ".outgroup.fa").is_file(): # WILL RAISE AN ERROR IF THIS EXISTS BECAUSE ALL_KEYS DOESNT in_keys, all_keys = outgroup_mngr( db, @@ -1847,14 +1828,13 @@ def search_main( verbose=False, ) query_hits = all_keys - out_query = query + ".outgroup" out_keys = list(set(all_keys).difference(set(in_keys))) - print("\t\t\t" + str(len(in_keys)) + " gene ingroup", flush=True) + logger.debug("" + str(len(in_keys)) + " gene ingroup") if out_keys: - print("\t\t\t" + str(len(out_keys)) + " gene outgroup", flush=True) + logger.debug("" + str(len(out_keys)) + " gene outgroup") else: query_fa = fa2dict(wrk_dir + query + ".fa") - null = crap_mngr( + crap_mngr( db, query, query_hits, @@ -1878,7 +1858,7 @@ def search_main( ) if output_loci: - print("\nOutputting most similar loci to query", flush=True) + logger.info("Outputting most similar loci to query") locus_output_mngr( gff_dir, loc_dir, @@ -1903,7 +1883,7 @@ def cli(): "-" for stdin', required=True, ) - i_opt.add_argument("-d", "--mtdb", default=primaryDB()) + i_opt.add_argument("-d", "--mtdb", default=primary_db()) i_opt.add_argument( "-g", "--gff", @@ -2044,37 +2024,38 @@ def cli(): help='Output format: ["svg", "pdf", "png"]', ) args = parser.parse_args() + setup_logging(verbose=getattr(args, "verbose", False)) execs = ["diamond", "clipkit", "mafft", "iqtree"] if args.search: if args.search not in {"mmseqs", "blastp", "diamond"}: - eprint("\nERROR: invalid -s", flush=True) + logger.error("invalid -s") sys.exit(3) else: execs.append(args.search) args.homologs = None - findExecs(execs, exit=set(execs)) + find_execs(execs, exit=set(execs)) if args.out_format.lower() not in {"svg", "pdf", "png"}: - eprint("\nERROR: invalid -of", flush=True) + logger.error("invalid -of") sys.exit(10) else: out_ext = "." + args.out_format.lower() if not args.homologs and args.faa: - eprint("\nERROR: -hf requires -hg", flush=True) + logger.error("-hf requires -hg") sys.exit(12) input_fa, input_GFF = False, False if args.query == "-": if args.gff: - eprint("\nERROR: GFF input requires fasta input", flush=True) + logger.error("GFF input requires fasta input") sys.exit(3) input_genes = stdin2str().replace('"', "").replace("'", "").split() elif "'" in args.query or '"' in args.query: if args.gff: - eprint("\nERROR: GFF input requires fasta input", flush=True) + logger.error("GFF input requires fasta input") sys.exit(3) input_genes = ( args.query.replace('"', "").replace("'", "").replace(",", " ").split() @@ -2085,10 +2066,9 @@ def cli(): ) db = mtdb(args.mtdb) - gene0 = input_genes[0] ome = input_genes[0][: input_genes[0].find("_")] if not ome in set(db["ome"]): - if os.path.isfile(args.query): + if Path(args.query).is_file(): if args.query.lower().endswith((".fasta", ".fa", ".faa", ".fna", ".fsa")): input_fa = fa2dict(args.query) input_genes = list(input_fa.keys()) @@ -2099,11 +2079,10 @@ def cli(): data = raw.read() input_genes = data.rstrip().split() else: - print("\nDetected non-mycotools input", flush=True) + logger.info("Detected non-mycotools input") if not input_fa or not args.search: - eprint( - "\tnon-mycotools input requires -s, -i as a fasta, optionally -g", - flush=True, + logger.info( + "non-mycotools input requires -s, -i as a fasta, optionally -g" ) sys.exit(4) @@ -2123,7 +2102,7 @@ def cli(): fast = True if args.agg_clus and args.linclust: - eprint("\nERROR: multiple clustering methods specified", flush=True) + logger.error("multiple clustering methods specified") sys.exit(5) elif args.agg_clus: clus_meth = "diamond" @@ -2174,7 +2153,7 @@ def cli(): args.plusminus, not args.no_label, ) - print("\nPreparing output directory", flush=True) + logger.info("Preparing output directory") out_dir, wrk_dir, gff_dir, tre_dir = make_output(output, new_log) hg_main( db, @@ -2210,7 +2189,7 @@ def cli(): args.plusminus, not args.no_label, ) - print("\nPreparing output directory", flush=True) + logger.info("Preparing output directory") out_dir, wrk_dir, gff_dir, tre_dir = make_output(output, new_log) search_main( db, diff --git a/mycotools/db2microsyntree.py b/mycotools/phylo/synteny.py similarity index 86% rename from mycotools/db2microsyntree.py rename to mycotools/phylo/synteny.py index 0142fda..e223e29 100755 --- a/mycotools/db2microsyntree.py +++ b/mycotools/phylo/synteny.py @@ -1,6 +1,6 @@ #!/usr/bin/env python3 -import os +import logging import sys import shutil import argparse @@ -9,19 +9,21 @@ import multiprocessing as mp from tqdm import tqdm from itertools import combinations -from collections import defaultdict, Counter -from mycotools.db2files import soft_main as symlink_files -from mycotools.db2hgs import id_near_schgs +from collections import defaultdict +from mycotools.mtdb.files import soft_main as symlink_files +from mycotools.cluster.db import id_near_schgs from mycotools.lib.kontools import ( format_path, - mkOutput, - findExecs, + mk_output, + find_execs, intro, - outro, - eprint, + setup_logging, ) -from mycotools.lib.dbtools import mtdb, primaryDB +from mycotools.lib.dbtools import mtdb, primary_db from mycotools.lib.biotools import gff2list +from pathlib import Path + +logger = logging.getLogger(__name__) def run_mmseqs( @@ -35,12 +37,12 @@ def run_mmseqs( ): symlink_files(["faa"], db, wrk_dir, verbose=False) # symlink proteomes cluster_res_file = wrk_dir + "homolog_groups.tsv" - if not os.path.isfile(cluster_res_file): # NEED to add to log removal + if not Path(cluster_res_file).is_file(): # NEED to add to log removal # be cautious about shell injection because we need to glob int(cpus) float(min_id) float(min_cov) - if not os.path.isdir(wrk_dir): + if not Path(wrk_dir).is_dir(): raise OSError("invalid working directory") elif not algorithm in {"mmseqs easy-linclust", "mmseqs easy-cluster"}: raise OSError("invalid mmseqs binary") @@ -71,18 +73,18 @@ def run_mmseqs( ) # stderr = subprocess.DEVNULL) shutil.move(wrk_dir + "cluster_cluster.tsv", cluster_res_file) - elif os.path.getsize(cluster_res_file): + elif Path(cluster_res_file).stat().st_size: mmseqs_cmd = 0 else: mmseqs_cmd = 1 if mmseqs_cmd: - eprint("\tERROR: cluster failed") + logger.error("cluster failed") sys.exit(1) - if os.path.isfile(wrk_dir + "cluster_all_seqs.fasta"): - os.remove(wrk_dir + "cluster_all_seqs.fasta") - if os.path.isfile(wrk_dir + "cluster_rep_seq.fasta"): - os.remove(wrk_dir + "cluster_rep_seq.fasta") - if os.path.isdir(wrk_dir + "tmp/"): + if Path(wrk_dir + "cluster_all_seqs.fasta").is_file(): + Path(wrk_dir + "cluster_all_seqs.fasta").unlink() + if Path(wrk_dir + "cluster_rep_seq.fasta").is_file(): + Path(wrk_dir + "cluster_rep_seq.fasta").unlink() + if Path(wrk_dir + "tmp/").is_dir(): shutil.rmtree(wrk_dir + "tmp/") return cluster_res_file @@ -198,12 +200,11 @@ def compile_cds(gff_list, ome, gene2hg): elif not prot: # if there isn't a valid accession it may mean the # mycotools curation did not work or the user did not curate # correctly - print(entry["attributes"], prot_prep_i0, prot_prep_i1) + logger.debug( + "%s %s %s", entry["attributes"], prot_prep_i0, prot_prep_i1 + ) if not fail: - print( - "\tWARNING: " + ome + " has proteins in gff with no Alias", - flush=True, - ) + logger.debug("" + ome + " has proteins in gff with no Alias") fail = True continue cds_dict[entry["seqid"]][prot].extend( @@ -233,9 +234,7 @@ def parse_loci(gff_path, ome, gene2hg, window=6): """obtain a set of tuples of HG pairs {(OG0, OG1)...}""" gff_list = gff2list(gff_path) # open here to improve pickling - hg_dict = compile_cds( - gff_list, os.path.basename(gff_path).replace(".gff3", ""), gene2hg - ) + hg_dict = compile_cds(gff_list, Path(gff_path).name.replace(".gff3", ""), gene2hg) pairs = [] for scaf, hgs in hg_dict.items(): # for each contig windows = [ @@ -265,14 +264,13 @@ def compile_loci(db, ome2i, gene2hg, window, cpus=1): def form_cooccur_array(cooccur_dict, ome2i): - count, hgx2i, size_dict, cooccur_arrays, i2hgx = 0, {}, {}, {}, {} + hgx2i, cooccur_arrays, i2hgx = {}, {}, {} cooccur_dict = { k: tuple(sorted(v)) for k, v in sorted(cooccur_dict.items(), key=lambda x: len(x[0])) } cooccur_arrays = np.zeros([len(ome2i), len(cooccur_dict)], dtype=np.int32) - old_len = len(list(cooccur_dict.keys())[0]) for i, hgx in enumerate(list(cooccur_dict.keys())): i2hgx[i] = hgx hgx2i[hgx] = i @@ -342,25 +340,23 @@ def align_microsynt_np(m_arr, i2ome, hg2gene, hgpair2i, wrk_dir, nschgs=None): ) if len(schgs) > 9: nschgs = schgs - print(f"\t\t{len(schgs)} HGs in single copy extracted", flush=True) + logger.debug(f"{len(schgs)} HGs in single copy extracted") elif len(nschgs) < 10: nschgs = [] if nschgs: - print( - f"\t\t{len(nschgs)} HGs with <= {max_median} median copy " - + f"number and <= {max_stdev} standard deviation extracted", - flush=True, + logger.debug( + f"{len(nschgs)} HGs with <= {max_median} median copy " + + f"number and <= {max_stdev} standard deviation extracted" ) max_stdev += 0.2 if max_stdev > 2: max_stdev = 0.1 max_median += 1 if not nschgs: - eprint( - "\nERROR: could not detect 10 genes present in all genomes " + logger.error( + "could not detect 10 genes present in all genomes " + "with median 2 copy number and less than 2 copy number " - + "standard deviation. Manually input focal homology groups.", - flush=True, + + "standard deviation. Manually input focal homology groups." ) sys.exit(35) @@ -395,8 +391,8 @@ def run_tree( ): tree_dir = wrk_dir + "tree/" - if not os.path.isdir(tree_dir): - os.mkdir(tree_dir) + if not Path(tree_dir).is_dir(): + Path(tree_dir).mkdir() prefix = tree_dir + "microsynt" tree_cmd = [iqtree, "-s", alignment, "-m", model, "--prefix", prefix] @@ -437,10 +433,10 @@ def main( # obtain useable omes useableOmes, dbOmes = set(), set(db.keys()) - print("\nI. Inputting data", flush=True) + logger.info("I. Inputting data") if n50thresh: # optional n50 threshold via mycotoolsDB assemblyPath = format_path("$MYCODB/../data/assemblyStats.tsv") - if os.path.isfile(assemblyPath): + if Path(assemblyPath).is_file(): with open(assemblyPath, "r") as raw: for line in raw: d = line.rstrip().split("\t") @@ -455,9 +451,9 @@ def main( useableOmes = dbOmes # initialize orthogroup data structures - if not hg_file and not os.path.isfile(wrk_dir + "homology_groups.tsv"): + if not hg_file and not Path(wrk_dir + "homology_groups.tsv").is_file(): hg_file = run_mmseqs(db, wrk_dir, algorithm=algorithm, min_id=min_id, cpus=cpus) - print("\tParsing homology groups (HGs)", flush=True) + logger.info("Parsing homology groups (HGs)") ome2i, gene2hg, i2ome, hg2gene = compile_homolog_groups( hg_file, wrk_dir, useableOmes ) @@ -468,9 +464,9 @@ def main( # remove genomes that are not in the db missing_from_db = set(ome2i.keys()).difference(set(db.keys())) - print("\t\tOmes:", len(ome2i), flush=True) + logger.info("%s %s", "\t\tOmes:", len(ome2i)) if missing_from_db: - print(f"\t\t\t{len(missing_from_db)} omes in HGs but not mtdb") + logger.debug(f"{len(missing_from_db)} omes in HGs but not mtdb") for ome in list(missing_from_db): del ome2i[ome] todel = [] @@ -489,32 +485,32 @@ def main( with open(wrk_dir + "ome2i.tsv", "w") as out: out.write("\n".join([k + "\t" + str(v) for k, v in ome2i.items()])) - print("\t\tHGs:", len(hg2gene), flush=True) - print("\t\tGenes:", len(gene2hg), flush=True) + logger.info("%s %s", "\t\tHGs:", len(hg2gene)) + logger.info("%s %s", "\t\tGenes:", len(gene2hg)) # compile cooccuring pairs of homogroups in each genome - print("\tCompiling all loci", flush=True) + logger.info("Compiling all loci") cc_arr_path = wrk_dir + "microsynt" ome2pairs = compile_loci(db, ome2i, gene2hg, plusminus * 2 + 1, cpus=cpus) cooccur_dict = None # if not os.path.isfile(out_dir + 'hgps.tsv.gz'): # assimilate cooccurrences across omes - print("\tIdentifying cooccurences", flush=True) + logger.info("Identifying cooccurences") seed_len = sum([len(ome2pairs[x]) for x in ome2pairs]) - print("\t\t" + str(seed_len) + " initial HG-pairs", flush=True) + logger.debug("" + str(seed_len) + " initial HG-pairs") cooccur_array, cooccur_dict, hgpair2i, i2hgpair = form_cooccur_structures( ome2pairs, 2, ome2i, cc_arr_path ) max_omes = max([len(cooccur_dict[x]) for x in cooccur_dict]) - print("\t\t" + str(max_omes) + " maximum organisms with HG-pair", flush=True) + logger.debug("" + str(max_omes) + " maximum organisms with HG-pair") cooccur_array[cooccur_array > 0] = 1 cooccur_array.astype(np.uint8) - print("\t\t" + str(sys.getsizeof(cooccur_array) / 1000000) + " MB", flush=True) + logger.debug("" + str(sys.getsizeof(cooccur_array) / 1000000) + " MB") cooccur_array, del_omes = remove_nulls(cooccur_array) for i in del_omes: - print(f"\t\t\t{i2ome[i]} removed for lacking overlap") + logger.debug(f"{i2ome[i]} removed for lacking overlap") del i2ome[i] ome2i = {v: i for i, v in enumerate(i2ome)} @@ -532,9 +528,8 @@ def main( # cooccur_array = np.load(cc_arr_path + '.npy') ome2pairs = {ome2i[ome]: v for ome, v in ome2pairs.items() if ome in ome2i} - microsynt_dict = {} - print("\nII. Microsynteny tree", flush=True) - if not os.path.isfile(tree_path): + logger.info("II. Microsynteny tree") + if not Path(tree_path).is_file(): nschgs = [] if near_single_copy_genes: try: @@ -550,11 +545,11 @@ def main( } ) # create microsynteny distance matrix and make tree - print("\tPreparing microsynteny alignment", flush=True) + logger.info("Preparing microsynteny alignment") align_file = align_microsynt_np( cooccur_array, i2ome, hg2gene, hgpair2i, wrk_dir, nschgs ) - print("\tBuilding microsynteny tree", flush=True) + logger.info("Building microsynteny tree") run_tree( align_file, wrk_dir, @@ -565,8 +560,8 @@ def main( cpus=cpus, ) # too bulky to justify keeping - if os.path.isfile(cc_arr_path + ".npy"): - os.remove(cc_arr_path + ".npy") + if Path(cc_arr_path + ".npy").is_file(): + Path(cc_arr_path + ".npy").unlink() return ome2i, gene2hg, i2ome, hg2gene, ome2pairs, cooccur_dict @@ -578,7 +573,7 @@ def cli(): + "loci." ) parser.add_argument( - "-d", "--db", default=primaryDB(), help="MycotoolsDB. DEFAULT: masterdb" + "-d", "--db", default=primary_db(), help="MycotoolsDB. DEFAULT: masterdb" ) parser.add_argument( "-f", @@ -612,17 +607,18 @@ def cli(): parser.add_argument("-c", "--cpus", default=1, type=int) parser.add_argument("-o", "--output") args = parser.parse_args() + setup_logging(verbose=getattr(args, "verbose", False)) execs = ["iqtree"] if args.orthofinder: of_out = format_path(args.orthofinder) - if os.path.isdir(of_out): + if Path(of_out).is_dir(): homogroups = of_out + "/Orthogroups/Orthogroups.txt" hg_dir = of_out + "/Orthogroup_Sequences/" else: homogroups = of_out - hg_dir = os.path.dirname(of_out) + "../Orthogroup_Sequences/" - if not os.path.isfile(hg_dir + "OG0000000.fa"): + hg_dir = str(Path(of_out).parent) + "../Orthogroup_Sequences/" + if not Path(hg_dir + "OG0000000.fa").is_file(): hg_dir = None method = "orthofinder" elif args.input: @@ -639,12 +635,12 @@ def cli(): hg_dir = None homogroups = None execs.append("mmseqs") - findExecs(execs, exit=set(execs)) + find_execs(execs, exit=set(execs)) if not args.output: - out_dir = mkOutput(os.getcwd() + "/", "db2microsyntree") + out_dir = mk_output(str(Path.cwd()) + "/", "phylo_synteny") else: - out_dir = mkOutput(format_path(args.output), "db2microsyntree") + out_dir = mk_output(format_path(args.output), "phylo_synteny") args_dict = { "Database": args.db, @@ -658,8 +654,8 @@ def cli(): intro("db2microsyntree", args_dict, "Zachary Konkel") wrk_dir = out_dir + "working/" - if not os.path.isdir(wrk_dir): - os.mkdir(wrk_dir) + if not Path(wrk_dir).is_dir(): + Path(wrk_dir).mkdir() if args.focal_genes: with open(format_path(args.focal_genes), "r") as raw: diff --git a/mycotools/treetools.py b/mycotools/phylo/tools.py similarity index 83% rename from mycotools/treetools.py rename to mycotools/phylo/tools.py index a2d26d7..2c07d6f 100755 --- a/mycotools/treetools.py +++ b/mycotools/phylo/tools.py @@ -1,20 +1,23 @@ #! /usr/bin/env python3 -import os +import logging import re import sys import argparse from itertools import chain -from cogent3 import PhyloNode, load_tree -from mycotools.lib.kontools import format_path, split_input, eprint, vprint +from cogent3 import load_tree +from mycotools.lib.kontools import format_path, split_input, setup_logging from mycotools.lib.dbtools import mtdb +from pathlib import Path + +logger = logging.getLogger(__name__) def compile_tree(tree_path, root=[], verbose=False): """Compile a phylogeny from a path and conver the tip names to an index""" phylo = load_tree(tree_path) tips = set(phylo.get_tip_names()) - vprint(f"{len(tips)} tips on input", v=verbose, e=True) + logger.debug(f"{len(tips)} tips on input") if len(root) == 1: phylo = phylo.rooted_with_tip(root[0]) @@ -27,9 +30,7 @@ def compile_tree(tree_path, root=[], verbose=False): mrca_tip_len = min([v[1] for v in list(nodes.values())]) mrca_edge = [k for k, v in nodes.items() if v[1] == mrca_tip_len] if mrca_edge[0] == "root": - eprint( - f"WARNING: rooting with {root} does not change current tree", flush=True - ) + logger.warning(f"rooting with {root} does not change current tree") phylo = phylo.rooted_at(mrca_edge[0]) return phylo @@ -39,7 +40,7 @@ def prune_to_new(phylo, tips=[], spacer=""): """Prune missing tips from a rooted phylogeny""" missing = set(phylo.get_tip_names()).difference(tips) if missing: - eprint(f"{spacer}Removing {len(missing)} tips", flush=True) + logger.info(f"{spacer}Removing {len(missing)} tips") todel = [] for n in phylo.tips(): if n.name in missing: @@ -97,13 +98,14 @@ def cli(): help="Omit support values", ) args = parser.parse_args() + setup_logging(verbose=getattr(args, "verbose", False)) root = split_input(args.root) if args.tips: tips_path = format_path(args.tips) - if os.path.isfile(tips_path): - eprint("\nDetecting path input", flush=True) + if Path(tips_path).is_file(): + logger.info("Detecting path input") with open(tips_path, "r") as raw: tips = list(chain(*[line.rstrip().split() for line in raw])) else: @@ -114,8 +116,8 @@ def cli(): convert = {} if args.convert: convert_path = format_path(args.convert) - if os.path.isfile(convert_path): - eprint("\nDetecting path input", flush=True) + if Path(convert_path).is_file(): + logger.info("Detecting path input") with open(convert_path, "r") as raw: for line in raw: data = line.rstrip() @@ -123,7 +125,7 @@ def cli(): k2v = data.split() convert[k2v[0]] = k2v[1] else: - eprint("\nERROR: --convert must be a valid reference file", flush=True) + logger.error("--convert must be a valid reference file") if args.mtdb: db = mtdb(format_path(args.mtdb)) @@ -136,7 +138,7 @@ def cli(): format_path(args.input), db, tips, trim=args.prune, root=root, convert=convert ) - eprint(f"{len(phylo.get_tip_names())} tips on output", flush=True) + logger.info(f"{len(phylo.get_tip_names())} tips on output") p_str = phylo.get_newick(with_distances=True) if args.support: print(p_str, flush=True) diff --git a/mycotools/fa2tree.py b/mycotools/phylo/tree.py similarity index 82% rename from mycotools/fa2tree.py rename to mycotools/phylo/tree.py index b6c53e8..18f942b 100755 --- a/mycotools/fa2tree.py +++ b/mycotools/phylo/tree.py @@ -4,6 +4,7 @@ # NEED to ignore non fasta inputs # NEED to work as a standalone script +import logging import os import re import sys @@ -11,31 +12,27 @@ import argparse import subprocess import contextlib -import multiprocessing as mp from collections import defaultdict from mycotools.lib.kontools import ( - eprint, - vprint, collect_files, format_path, intro, outro, - findExecs, - mkOutput, + find_execs, + mk_output, multisub, - parse_run_log, + setup_logging, ) from mycotools.lib.biotools import fa2dict, dict2fa +from pathlib import Path + +logger = logging.getLogger(__name__) -try: - from clipkit import clipkit -except ImportError: - eprint("ERROR: clipkit is not installed. Install via `conda` or `pip`") try: from ete3 import Tree except ImportError: - eprint( - "WARNING: ete3 not installed.\nInstall ete3 into your " + logger.warning( + "ete3 not installed.\nInstall ete3 into your " + "conda environment via `conda install ete3`" ) @@ -68,7 +65,7 @@ def run_mafft( # run the command directly if not hpc: - print(spacer + "Aligning", flush=True) + logger.debug(spacer + "Aligning") with open(out_dir + name + ".mafft", "w") as out_file: if verbose: run_mafft = subprocess.call(cmd, stdout=out_file) @@ -78,9 +75,9 @@ def run_mafft( ) if run_mafft != 0: - eprint(spacer + "\tERROR: mafft failed: " + str(run_mafft), flush=True) - if os.path.isfile(out_dir + name + ".mafft"): - os.remove(out_dir + name + ".mafft") + logger.error(spacer + "mafft failed: " + str(run_mafft)) + if Path(out_dir + name + ".mafft").is_file(): + Path(out_dir + name + ".mafft").unlink() if pass_fail: return None raise PhyloError @@ -128,9 +125,8 @@ def run_clipkit( ): """Execute ClipKIT from a complete Mafft run""" - clipkit_out_name = out_dir + os.path.basename(mafft_name) + ".clipkit" + clipkit_out_name = out_dir + Path(mafft_name).name + ".clipkit" if gappy: - mode = "gappy" cmd = [ "clipkit", mafft_name, @@ -142,44 +138,21 @@ def run_clipkit( ] else: cmd = ["clipkit", mafft_name, "--output", clipkit_out_name] - mode = "smart-gap" gappy = None # execute immediately if not hpc: - print(spacer + "Trimming", flush=True) + logger.debug(spacer + "Trimming") if not verbose: clipkit_code = subprocess.call(cmd, stdout=subprocess.PIPE) else: clipkit_code = subprocess.call( cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL ) - # clipkit_out = clipkit.execute( - # input_file=mafft_name, - # output_file=clipkit_out_name, - # output_file_format='fasta', - # input_file_format='fasta', - # use_log = False, - # complement = False, - # mode=mode, gaps=gappy - # ) - # else: - # with nostdout(): - # clipkit_out = clipkit.execute( - # input_file=mafft_name, - # output_file=clipkit_out_name, - # output_file_format='fasta', - # input_file_format='fasta', - # use_log = False, - # complement = False, - # mode=mode, gaps=gappy - # ) # no output file, the run failed - if not os.path.isfile(clipkit_out_name): - eprint( - spacer + "\tERROR: `clipkit` failed: " + str(clipkit_code), flush=True - ) + if not Path(clipkit_out_name).is_file(): + logger.error(spacer + "`clipkit` failed: " + str(clipkit_code)) if pass_fail: return None raise PhyloError @@ -227,7 +200,6 @@ def run_mf( """Run ModelFinder on its own in preparation for multigene phylogeny reconstruction""" # set the CPUs for each ModelFinder run arbitrarily - cpus_per_cmd = 3 # determine the concurrent ModelFinder runs that are possible concurrent_cmds = round((cpus - 1) / 4 - 0.5) # round down if not concurrent_cmds: @@ -246,10 +218,10 @@ def run_mf( "-s", f_, "--prefix", - out_dir + os.path.basename(f_), + out_dir + Path(f_).name, ] for f_ in clipkit_files - if not os.path.isfile(f"{out_dir}{os.path.basename(f_)}.contree") + if not Path(f"{out_dir}{Path(f_).name}.contree").is_file() ] # append the topological constraint to each command if it is present if constraint: @@ -260,9 +232,9 @@ def run_mf( clipkit_files = [ x for x in clipkit_files - if not os.path.isfile(f"{out_dir}{os.path.basename(x)}.contree") + if not Path(f"{out_dir}{Path(x).name}.contree").is_file() ] - return {os.path.basename(v): cmds[i] for i, v in enumerate(clipkit_files)} + return {Path(v).name: cmds[i] for i, v in enumerate(clipkit_files)} # otherwise parallelize and run ModelFinder directly multisub(cmds, verbose=verbose, processes=concurrent_cmds) @@ -270,7 +242,7 @@ def run_mf( # parse and identify the evolutionary models determined by ModelFinder models = {} for f_ in clipkit_files: - with open(out_dir + os.path.basename(f_) + ".log", "r") as raw: + with open(out_dir + Path(f_).name + ".log", "r") as raw: for line in raw: if line.startswith("Best-fit model:"): model = re.search(r"Best-fit model: (.+) chosen", line)[1] @@ -293,11 +265,8 @@ def prepare_nexus(concat_fa, models, spacer="\t"): # if some alignments aren't the same length then there is some cryptic # issue, likely user-caused if not all(len(x["sequence"]) == len0 for x in trim_fa.values()): - eprint( - spacer - + "\tERROR: alignment sequences are not same length " - + trimmed_f, - flush=True, + logger.error( + spacer + "alignment sequences are not same length " + trimmed_f ) sys.exit(17) # adjust the index of the coordinates of each sequence based on the @@ -356,11 +325,11 @@ def run_partition_tree( if constraint: cmd.extend(["-g", constraint]) if tree_stop: - print(spacer + "Concatenated nexus and fasta outputted", flush=True) - print(" ".join(cmd), flush=True) + logger.debug(spacer + "Concatenated nexus and fasta outputted") + logger.debug(" ".join(cmd)) sys.exit(0) - print(spacer + "Tree building", flush=True) + logger.debug(spacer + "Tree building") if verbose: run_tree = subprocess.call(cmd) else: @@ -384,7 +353,7 @@ def run_tree_reconstruction( """Manage and execute phylogeny reconstruction from an inputted ClipKIT output trimmed alignment.""" - tree_file = out_dir + os.path.basename(clipkit_file) + tree_file = out_dir + Path(clipkit_file).name # prepare a fasttree ommand if fast: cmd = ["fasttree", "-out", tree_file + ".treefile", clipkit_file] @@ -409,7 +378,7 @@ def run_tree_reconstruction( # execute in the current terminal if not hpc: - print(spacer + "Tree building", flush=True) + logger.debug(spacer + "Tree building") if fast: cmd[2] += ".tmp" if verbose: @@ -421,12 +390,12 @@ def run_tree_reconstruction( if fast: shutil.move(cmd[2], cmd[2][:-4]) if run_tree != 0: - eprint(spacer + "\tERROR: tree failed: " + str(run_tree), flush=True) + logger.error(spacer + "tree failed: " + str(run_tree)) raise PhyloError # prepare an .sh file for user execution, or sequential execution from the # previous scripts else: - vprint("\nOutputting bash script `tree.sh`.\n", v=verbose, flush=True) + logger.debug("Outputting bash script `tree.sh`.\n") with open(f"{out_dir}../{prefix}_tree.sh", "w") as out: out.write(hpc + "\n\n" + " ".join([str(x) for x in cmd])) @@ -435,37 +404,37 @@ def prep_fasta_path_input(fasta_path, output_dir): """Prepare the output directories and collect the files from an input that is a fasta path (file or directory)""" # start from a file input - if os.path.isfile(fasta_path): + if Path(fasta_path).is_file(): if not output_dir: dir_name = re.sub( - r"\..*?$", "_tree", os.path.basename(os.path.abspath(fasta_path)) + r"\..*?$", "_tree", Path(os.path.abspath(fasta_path)).name ) out_dir = output_dir + "/" + dir_name - if not os.path.isdir(out_dir): - os.mkdir(out_dir) + if not Path(out_dir).is_dir(): + Path(out_dir).mkdir() out_dir = format_path(out_dir) else: - if not os.path.isdir(output_dir): - os.mkdir(output_dir) + if not Path(output_dir).is_dir(): + Path(output_dir).mkdir() out_dir = format_path(output_dir) wrk_dir = out_dir + "working/" - if not os.path.isdir(wrk_dir): - os.mkdir(wrk_dir) + if not Path(wrk_dir).is_dir(): + Path(wrk_dir).mkdir() files = [fasta_path] # start from a directory of fastas - elif os.path.isdir(fasta_path): + elif Path(fasta_path).is_dir(): if not output_dir: - out_dir = mkOutput(output_dir_prep, "fa2tree") + out_dir = mk_output(output_dir_prep, "phylo_tree") else: - if not os.path.isdir(output_dir): - os.mkdir(output_dir) + if not Path(output_dir).is_dir(): + Path(output_dir).mkdir() out_dir = format_path(output_dir) wrk_dir = out_dir + "working/" - if not os.path.isdir(wrk_dir): - os.mkdir(wrk_dir) + if not Path(wrk_dir).is_dir(): + Path(wrk_dir).mkdir() files = collect_files(fasta_path, "*") check_fas = set(files) @@ -475,11 +444,11 @@ def prep_fasta_path_input(fasta_path, output_dir): if f + ".mafft" not in check_fas: new_fas.append(f) else: - shutil.copy(f + ".mafft", wrk_dir + os.path.basename(f) + ".mafft") + shutil.copy(f + ".mafft", wrk_dir + Path(f).name + ".mafft") else: shutil.copy( f + ".mafft.clipkit", - wrk_dir + os.path.basename(f) + ".mafft.clipkit", + wrk_dir + Path(f).name + ".mafft.clipkit", ) files = new_fas return out_dir, wrk_dir, files @@ -487,17 +456,17 @@ def prep_fasta_path_input(fasta_path, output_dir): def prep_fasta_list_input(fastas, output_dir): if not output_dir: - out_dir = mkOutput("./", "fa2tree") + out_dir = mk_output("./", "phylo_tree") else: - if not os.path.isdir(output_dir): - os.mkdir(output_dir) + if not Path(output_dir).is_dir(): + Path(output_dir).mkdir() out_dir = format_path(output_dir) wrk_dir = out_dir + "working/" - if not os.path.isdir(wrk_dir): - os.mkdir(wrk_dir) + if not Path(wrk_dir).is_dir(): + Path(wrk_dir).mkdir() files = [] for path in fastas: - if os.path.isdir(path): + if Path(path).is_dir(): files.extend([format_path(x) for x in collect_files(path, "*")]) else: files.append(path) @@ -509,10 +478,8 @@ def nonfasta2fasta(files, conv_dir): from Bio import SeqIO for i, f_ in enumerate(files): - out_name = conv_dir + os.path.basename( - re.search(r"(.*)\.[^.]+$", f_)[1] + ".fa" - ) - if not os.path.isfile(out_name): + out_name = conv_dir + Path(re.search(r"(.*)\.[^.]+$", f_)[1] + ".fa").name + if not Path(out_name).is_file(): if f_.endswith( ( ".nexus", @@ -555,21 +522,18 @@ def identify_incomplete_files(files, flag_incomplete, wrk_dir): if f.endswith("/"): f = f[:-1] fa_dict = fa2dict(f) - incomp_omes[os.path.basename(f)] = [] + incomp_omes[Path(f).name] = [] # populate a hash with each gene for each ome associated with each # fasta for seq in fa_dict: ome = seq[: seq.find("_")] if f not in ome2fa2gene[ome]: - incomp_omes[os.path.basename(f)].append(ome) + incomp_omes[Path(f).name].append(ome) ome2fa2gene[ome][f] = seq # a multigene partition cannot be reconstructed for a genome with # multiple genes in the same alignment else: - eprint( - "\nERROR: multiple sequences for a " + "single ome in " + f, - flush=True, - ) + logger.error("multiple sequences for a " + "single ome in " + f) sys.exit(5) # identify the files with missing genomes and the missing omes themselves @@ -591,23 +555,21 @@ def identify_incomplete_files(files, flag_incomplete, wrk_dir): # it is an error if genomes are missing and it isn't explicitly # permitted if flag_incomplete: - eprint( - "\nERROR: omes without sequences in all fastas\n", - +"Run with -m to remove failed omes", - flush=True, + logger.error( + "omes without sequences in all fastas\n" + "Run with -m to remove failed omes" ) else: - eprint( - "\nWARNING: omes removed without sequences in all \ - fastas: ", - flush=True, + logger.warning( + "omes removed without sequences in all \ + fastas: " ) for f, omes in incomp_files.items(): - eprint("\t" + f + ": " + ",".join([x for x in omes]), flush=True) + logger.info("" + f + ": " + ",".join([x for x in omes])) if comp_files: - eprint("\tComplete files: " + ",".join(comp_files), flush=True) + logger.info("Complete files: " + ",".join(comp_files)) else: - eprint("\tNo complete files", flush=True) + logger.info("No complete files") # exit if incomplete flags if flag_incomplete: sys.exit(6) @@ -616,14 +578,14 @@ def identify_incomplete_files(files, flag_incomplete, wrk_dir): else: new_files = [] for fasta in files: - fasta_name = os.path.basename(fasta) + fasta_name = Path(fasta).name fa = fa2dict(fasta) new_fa = { k: v for k, v in fa.items() if k[: k.find("_")] not in del_omes } with open(wrk_dir + fasta_name + ".tmp", "w") as out: out.write(dict2fa(new_fa)) - os.rename(wrk_dir + fasta_name + ".tmp", wrk_dir + fasta_name) + Path(wrk_dir + fasta_name + ".tmp").rename(wrk_dir + fasta_name) new_files.append(wrk_dir + fasta_name) files = new_files @@ -638,7 +600,7 @@ def algn_mngr( trimmed_files = [] for fasta in files: name = re.search(r".*?/*([^/]+)/*$", fasta)[1] - fasta_name = os.path.basename(os.path.abspath(fasta)) + fasta_name = Path(os.path.abspath(fasta)).name clipkit = wrk_dir + fasta_name + ".clipkit" # if the starting point is a non-aligned fasta @@ -647,11 +609,11 @@ def algn_mngr( clipkit = mafft + ".clipkit" try: # check if it exists - if os.stat(mafft): - vprint("\nAlignment exists", v=verbose, flush=True) + if Path(mafft).stat(): + logger.debug("Alignment exists") else: raise ValueError - except (FileNotFoundError, ValueError) as e: + except (FileNotFoundError, ValueError): # otherwise run the alignment if not alignment: mafft = run_mafft( @@ -677,12 +639,12 @@ def algn_mngr( if start == 0 or start == 1: try: # check for a trimmed fasta - if os.stat(clipkit): - vprint("\nTrim exists", v=verbose, flush=True) + if Path(clipkit).stat(): + logger.debug("Trim exists") clipkit_out = clipkit else: raise ValueError - except (FileNotFoundError, ValueError) as e: + except (FileNotFoundError, ValueError): clipkit_out = run_clipkit( name, mafft, @@ -706,7 +668,7 @@ def convert_seq_to_ome_name(trimmed_files, conv_dir): """Convert the inputted sequence name to its genome name""" new_trimmed_files = [] for trimmed_f in trimmed_files: - new_f = conv_dir + os.path.basename(trimmed_f) + new_f = conv_dir + Path(trimmed_f).name in_fa = fa2dict(trimmed_f) out_fa = {} for seq, data in in_fa.items(): @@ -724,19 +686,18 @@ def extract_supported(trimmed_files, min_mean_support, out_dir): """Extract trees that meet the minimum summary support value""" out_files = [] for f_ in trimmed_files: - t_path = f"{out_dir}working/{os.path.basename(f_)}.contree" + t_path = f"{out_dir}working/{Path(f_).name}.contree" supports = check_tree_support(t_path) if supports is None: continue mean_support = sum(supports) / len(supports) if mean_support >= min_mean_support: - print(f"\t{os.path.basename(t_path)} {mean_support} passed", flush=True) + logger.debug(f"{Path(t_path).name} {mean_support} passed") out_files.append(f_) else: - print(f"\t{os.path.basename(t_path)} {mean_support} failed", flush=True) - print( - f"\t{len(out_files)} ({len(out_files)/len(trimmed_files)*100}%)" + " passed", - flush=True, + logger.debug(f"{Path(t_path).name} {mean_support} failed") + logger.debug( + f"{len(out_files)} ({len(out_files)/len(trimmed_files)*100}%)" + " passed" ) return out_files @@ -781,21 +742,21 @@ def multigene_mngr( + f"#SBATCH --job-name={f}_tree\n" ) out.write(" ".join(cmd)) - print(spacer + "Alignments outputed", flush=True) + logger.debug(spacer + "Alignments outputed") sys.exit(0) - print(spacer + "Model finding", flush=True) + logger.debug(spacer + "Model finding") models = run_mf( trimmed_files, wrk_dir, constraint, verbose, cpus, spacer=spacer + "\t" ) - print(f"{spacer}\t{len(models)} individual trees completed", flush=True) + logger.debug(f"{spacer}\t{len(models)} individual trees completed") if min_mean_support: - print(spacer + "Extracting passing trees", flush=True) + logger.debug(spacer + "Extracting passing trees") passing = extract_supported(trimmed_files, min_mean_support, out_dir) models = {k: models[k] for k in passing} # build the concatenated NEXUS - print(spacer + "Concatenating", flush=True) + logger.debug(spacer + "Concatenating") empty_concat_fa = { ome: {"description": "", "sequence": ""} for ome in ome2fa2gene @@ -872,9 +833,7 @@ def main( # prepare the scaffold of an HPC command if that is requested hpc, hpc_prep = False, False if slurm: - vprint( - "\nHPC mode, preparing submission scripts (Slurm)\n", v=verbose, flush=True - ) + logger.debug("HPC mode, preparing submission scripts (Slurm)\n") hpc_prep = ( "#!/bin/bash\n#SBATCH --time=24:00:00\n" + "#SBATCH --nodes=1\n#SBATCH --ntasks-per-node=" @@ -883,16 +842,13 @@ def main( ) # '\n\nsource activate ' + source elif torque: - vprint( - "\nHPC mode, preparing submission scripts (PBS)\n", v=verbose, flush=True - ) + logger.debug("HPC mode, preparing submission scripts (PBS)\n") hpc_prep = ( "#PBS -l walltime=10:00:00\n#PBS -l nodes=1:ppn=4\n#PBS " + "-A " + project ) if not partition and hpc_prep: hpc = hpc_prep - output_dir_prep = os.getcwd() + "/" # the fasta data should be a path if it is a string if isinstance(fasta_path, str): out_dir, wrk_dir, files = prep_fasta_path_input(fasta_path, output_dir) @@ -902,8 +858,8 @@ def main( # create a directory for converting files conv_dir = wrk_dir + "conv/" - if not os.path.isdir(conv_dir): - os.mkdir(conv_dir) + if not Path(conv_dir).is_dir(): + Path(conv_dir).mkdir() # check for non fasta inputs - if these files exist in a directory then # they will be used anyway @@ -922,7 +878,7 @@ def main( ) for f_ in files ): - print(spacer + "Converting to fastas", flush=True) + logger.debug(spacer + "Converting to fastas") files = nonfasta2fasta(files, conv_dir) # only proceed with fastas from the inputted files @@ -978,11 +934,12 @@ def main( if hpc: if slurm: - print("\nStart pipeline via `sbatch .sh` in " + out_dir + "\n") + logger.debug( + "Start pipeline via `sbatch .sh` in " + out_dir + "\n" + ) else: - vprint( - "\nStart pipeline via `qsub .sh` in " + out_dir + "\n", - v=verbose, + logger.debug( + "Start pipeline via `qsub .sh` in " + out_dir + "\n" ) @@ -1067,6 +1024,7 @@ def cli(): r_opt.add_argument("-c", "--cpus", default=1, type=int) args = parser.parse_args() + setup_logging(verbose=getattr(args, "verbose", False)) execs = ["mafft", "clipkit"] if args.fast: @@ -1074,13 +1032,13 @@ def cli(): else: execs.append("iqtree") if args.torque or args.slurm: - findExecs(execs, execs) + find_execs(execs, execs) else: - findExecs(execs, execs) + find_execs(execs, execs) if args.support: if args.support > 1 or args.support < 0: - eprint("\nERROR: --support must be between 0 and 1") + logger.error("--support must be between 0 and 1") sys.exit(3) output = format_path(args.output) @@ -1090,7 +1048,7 @@ def cli(): if args.gappy: if args.gappy > 1: - eprint("\nERROR: gappy threshold must be less than 1") + logger.error("gappy threshold must be less than 1") sys.exit(3) args_dict = { diff --git a/mycotools/ome2name.py b/mycotools/rename.py similarity index 90% rename from mycotools/ome2name.py rename to mycotools/rename.py index d50b9ea..1c72c3b 100755 --- a/mycotools/ome2name.py +++ b/mycotools/rename.py @@ -1,10 +1,13 @@ #! /usr/bin/env python3 -import os import re import sys -from mycotools.lib.kontools import format_path, sys_start, eprint -from mycotools.lib.dbtools import primaryDB, mtdb +import logging +from mycotools.lib.kontools import format_path, sys_start, setup_logging +from mycotools.lib.dbtools import primary_db, mtdb +from pathlib import Path + +logger = logging.getLogger(__name__) def parse_args(args): @@ -19,7 +22,6 @@ def parse_args(args): ) allowable = {"_", "-", "!", "`", ",", ".", "~", "'", '"'} - arg_index = len(args) - 1 valid_ranks = {"kingdom", "subphylum", "phylum", "class", "order", "family"} tax, rank = False, None for arg in args[2:]: @@ -32,7 +34,7 @@ def parse_args(args): if tax: rank = arg.lower() if rank not in valid_ranks: - eprint("\nERROR: --taxonomy not in " + str(valid_ranks), flush=True) + logger.error("--taxonomy not in " + str(valid_ranks)) sys.exit(10) tax = False elif "-" in arg: @@ -40,11 +42,11 @@ def parse_args(args): tax = True continue else: - eprint("\nERROR: invalid argument `-`", flush=True) + logger.error("invalid argument `-`") sys.exit(11) - elif os.path.isfile(format_path(arg)): + elif Path(format_path(arg)).is_file(): if not go_on: - eprint("\nERROR: multiple files", flush=True) + logger.error("multiple files") db = mtdb(arg) go_on = False continue @@ -65,7 +67,7 @@ def parse_args(args): # import primary MTDB if go_on: - db = mtdb(primaryDB()) + db = mtdb(primary_db()) # read input data with open(args[1], "r") as raw_input: @@ -157,6 +159,7 @@ def main( def cli(): """Command line entry point""" + setup_logging() usage = ( "USAGE: ome2name | ome2name.py " + " [.mtdb] asvg*&\nDEFAULTS: master db, see script for default" diff --git a/mycotools/seq/__init__.py b/mycotools/seq/__init__.py new file mode 100644 index 0000000..18e15e9 --- /dev/null +++ b/mycotools/seq/__init__.py @@ -0,0 +1,42 @@ +#! /usr/bin/env python3 +"""Dispatcher for the `mycotools seq` subcommand. + +Routes `mycotools seq ...` to a sequence/coordinate transform module.""" +from mycotools.lib.subcmd import Dispatcher + +# subcommand name/alias -> submodule within this package (mycotools.seq.) +SUBCOMMANDS = { + "translate": "translate", + "coords": "coords", + "gff": "gff", + "convert": "convert", + "mass": "mass", +} + +DESCRIPTION = """Sequence and coordinate transforms + +Tools (all following arguments are forwarded to the tool): + translate translate a nucleotide fasta to protein + coords extract subsequence(s) from a fasta by coordinate + gff extract sequences from a gff3 (+ optional assembly) + convert convert between sequence file formats + mass compute protein masses from an amino-acid fasta + +Examples: + mycotools seq translate -h + mycotools seq coords -h""" + +_dispatcher = Dispatcher( + "mycotools seq", + "mycotools.seq", + SUBCOMMANDS, + DESCRIPTION, + metavar="TOOL", + arg_help="sequence tool (see below)", +) +main = _dispatcher.main +cli = _dispatcher.cli + + +if __name__ == "__main__": + cli() diff --git a/mycotools/seq/__main__.py b/mycotools/seq/__main__.py new file mode 100644 index 0000000..f51e01d --- /dev/null +++ b/mycotools/seq/__main__.py @@ -0,0 +1,6 @@ +#! /usr/bin/env python3 +"""Enable ``python -m mycotools.seq`` to run the seq dispatcher.""" +from mycotools.seq import cli + +if __name__ == "__main__": + cli() diff --git a/mycotools/bioreform.py b/mycotools/seq/convert.py similarity index 100% rename from mycotools/bioreform.py rename to mycotools/seq/convert.py diff --git a/mycotools/coords2fa.py b/mycotools/seq/coords.py similarity index 86% rename from mycotools/coords2fa.py rename to mycotools/seq/coords.py index d266ee9..2d75eb2 100755 --- a/mycotools/coords2fa.py +++ b/mycotools/seq/coords.py @@ -2,16 +2,18 @@ # NEED to flag overlapping coordinates -import os +import logging import sys -import argparse from Bio.Seq import Seq from collections import defaultdict from mycotools.lib.biotools import fa2dict, dict2fa -from mycotools.lib.kontools import sys_start, eprint, format_path +from mycotools.lib.kontools import sys_start, format_path +from pathlib import Path +logger = logging.getLogger(__name__) -def extractCoords(fa_dict, seqid, coord_start=0, coord_end=-1, sense="+", fa_name=""): + +def extract_coords(fa_dict, seqid, coord_start=0, coord_end=-1, sense="+", fa_name=""): """Extract the coordinates of a fasta based on input parameters""" new_fa, error = {}, "" @@ -67,9 +69,9 @@ def cli(): try: fa_file = format_path(args[0]) if fa_file.endswith("/"): # if it is a directory that was inputted - fa_name = os.path.basename(fa_file[:-1]) + fa_name = Path(fa_file[:-1]).name else: - fa_name = os.path.basename(fa_file) + fa_name = Path(fa_file).name fa = fa2dict(fa_file) @@ -80,7 +82,7 @@ def cli(): args.extend([0, -1, "+"]) elif len(args) < 6: args.append("+") - out_fa, error = extractCoords( + out_fa, error = extract_coords( fa, args[1], min([int(args[2]), int(args[3])]), @@ -91,7 +93,7 @@ def cli(): # print the extracted sequences to stdout in FASTA format print(dict2fa(out_fa)) - eprint(error) + logger.error(error) sys.exit(0) except IndexError: # fasta was not parseable @@ -125,7 +127,7 @@ def cli(): ) except: - eprint("\nERROR: incorrectly formatted input", flush=True) + logger.error("incorrectly formatted input") # for each file and coordinates, prepare for output for fa_file, concats in files_data.items(): @@ -133,9 +135,9 @@ def cli(): # extract the file name from the path if fa_file.endswith("/"): - fa_name = os.path.basename(fa_file[:-1]) + fa_name = Path(fa_file[:-1]).name else: - fa_name = os.path.basename(fa_file) + fa_name = Path(fa_file).name for concat_id, rows in concats.items(): if rows[0][-1] == "+": @@ -144,7 +146,9 @@ def cli(): sorted_rows = sorted(rows, key=lambda x: x[1], reverse=True) if concat_id is None: for x in sorted_rows: - new_fa, error_t = extractCoords(fa, x[0], x[1], x[2], x[3], fa_name) + new_fa, error_t = extract_coords( + fa, x[0], x[1], x[2], x[3], fa_name + ) error += error_t out_fa = { **out_fa, @@ -159,14 +163,16 @@ def cli(): ) toadd_fa = {"description": "", "sequence": ""} for x in sorted_rows: - new_fa, error_t = extractCoords(fa, x[0], x[1], x[2], x[3], fa_name) + new_fa, error_t = extract_coords( + fa, x[0], x[1], x[2], x[3], fa_name + ) error += error_t for seq, seq_info in new_fa.items(): toadd_fa["sequence"] += seq_info["sequence"] out_fa = {**out_fa, **{fa_name + "_concat" + str(concat_id): toadd_fa}} print(dict2fa(out_fa)) - eprint(error) + logger.error(error) sys.exit(0) diff --git a/mycotools/gff2seq.py b/mycotools/seq/gff.py similarity index 84% rename from mycotools/gff2seq.py rename to mycotools/seq/gff.py index c230d2b..4bb4852 100755 --- a/mycotools/gff2seq.py +++ b/mycotools/seq/gff.py @@ -1,15 +1,18 @@ #! /usr/bin/env python3 +import logging import re import sys import argparse from Bio.Seq import Seq -from mycotools.lib.dbtools import mtdb, primaryDB -from mycotools.lib.biotools import fa2dict, gff2list, gff3Comps, dict2fa -from mycotools.lib.kontools import format_path, sys_start, eprint, stdin2str +from mycotools.lib.dbtools import mtdb, primary_db +from mycotools.lib.biotools import fa2dict, gff2list, gff3_comps, dict2fa +from mycotools.lib.kontools import format_path, stdin2str, setup_logging +logger = logging.getLogger(__name__) -def sortGene(sorting_group): + +def sort_gene(sorting_group): out_group = [] for entryType in ["gene", "mrna", "trna", "rrna", "exon", "cds"]: @@ -27,7 +30,7 @@ def sortGene(sorting_group): return out_group -def sortContig(contigData): +def sort_contig(contigData): coordinates = {} for gene in contigData: @@ -45,9 +48,9 @@ def sortContig(contigData): return outContig -def sortGFF(unsorted_gff, idComp): +def sort_gff(unsorted_gff, idComp): - sorting_groups, oldGene = {}, None + sorting_groups = {} for i, entry in enumerate(unsorted_gff): seqid = entry["seqid"] if seqid not in sorting_groups: @@ -62,23 +65,21 @@ def sortGFF(unsorted_gff, idComp): for seqid in sorting_groups: contigData = {} for gene in sorting_groups[seqid]: - contigData[gene] = sortGene(sorting_groups[seqid][gene]) - sortedGff.extend(sortContig(contigData)) + contigData[gene] = sort_gene(sorting_groups[seqid][gene]) + sortedGff.extend(sort_contig(contigData)) return sortedGff -def sortMain(gff): +def sort_main(gff): - id_comp = re.compile(gff3Comps()["id"]) - # crude_sort = sorted( gff, key = lambda x: \ - # int( re.search(r'ID=([^;]+)', x['attributes'])[1] )) - gff = sortGFF(gff, id_comp) + id_comp = re.compile(gff3_comps()["id"]) + gff = sort_gff(gff, id_comp) return gff -def grabCDS(gff_dicts, spacer="\t"): +def grab_cds(gff_dicts, spacer="\t"): """Grab CDSs that are associated with genes. gff_dicts is a mycotools.lib.biotools gff2list() list""" @@ -86,7 +87,7 @@ def grabCDS(gff_dicts, spacer="\t"): for entry in gff_dicts: if entry["type"] == "mRNA": try: - alias = re.search(gff3Comps()["Alias"], entry["attributes"])[1] + alias = re.search(gff3_comps()["Alias"], entry["attributes"])[1] if not ome: ome = alias[: alias.find("_")] except TypeError: @@ -94,21 +95,20 @@ def grabCDS(gff_dicts, spacer="\t"): mrnas.extend(alias.split("|")) # account for posttranslational mods elif "gene" in entry["type"]: try: - alias = re.search(gff3Comps()["Alias"], entry["attributes"])[1] + alias = re.search(gff3_comps()["Alias"], entry["attributes"])[1] genes.extend(alias.split("|")) except TypeError: if not warning and ome: - eprint( - spacer + "WARNING: " + str(ome) + " missing aliases", flush=True - ) + logger.warning(spacer + "" + str(ome) + " missing aliases") warning = True raise TypeError(str(entry)) - mrna_set = set([x for x in mrnas if x in set(genes)]) + gene_set = set(genes) + mrna_set = set(x for x in mrnas if x in gene_set) out_cds = [] for entry in gff_dicts: if entry["type"] == "CDS": - alias = re.search(gff3Comps()["Alias"], entry["attributes"])[1] + alias = re.search(gff3_comps()["Alias"], entry["attributes"])[1] if alias in mrna_set: out_cds.append(entry) @@ -128,11 +128,11 @@ def order_neg_dict(neg_dict): return neg_dict -def grabCoords(cdss): +def grab_coords(cdss): pos_dict, neg_dict = {}, {} for entry in cdss: - gene = re.search(gff3Comps()["Alias"], entry["attributes"])[1] + gene = re.search(gff3_comps()["Alias"], entry["attributes"])[1] seqid = entry["seqid"] if entry["strand"].rstrip() == "+": if seqid not in pos_dict: @@ -141,7 +141,6 @@ def grabCoords(cdss): pos_dict[seqid][gene] = [] pos_dict[seqid][gene].append([int(entry["start"]) - 1, int(entry["end"])]) elif entry["strand"].rstrip() == "-": - negseq = seqid if seqid not in neg_dict: neg_dict[seqid], neg_dict[seqid][gene] = {}, [] elif gene not in neg_dict[seqid]: @@ -151,7 +150,7 @@ def grabCoords(cdss): return pos_dict, order_neg_dict(neg_dict) -def translatePos(contig_seq, postig_dict): +def translate_pos(contig_seq, postig_dict): genes_fa_dict = {} for gene in postig_dict: @@ -164,7 +163,7 @@ def translatePos(contig_seq, postig_dict): return genes_fa_dict -def translateNeg(rev_seq, negtig_dict): +def translate_neg(rev_seq, negtig_dict): genes_fa_dict = {} seqlen = len(rev_seq.rstrip()) @@ -180,7 +179,7 @@ def translateNeg(rev_seq, negtig_dict): return genes_fa_dict -def ntPos(contig_seq, postig_dict): +def nt_pos(contig_seq, postig_dict): genes_fa_dict = {} for gene in postig_dict: @@ -192,7 +191,7 @@ def ntPos(contig_seq, postig_dict): return genes_fa_dict -def posPlusMinusCode(genePostig_dict, plusminus, contig_seq, plus=True, minus=True): +def pos_plus_minus_code(genePostig_dict, plusminus, contig_seq, plus=True, minus=True): geneDict = {"sequence": "", "description": ""} minimumList = [] for x in genePostig_dict: @@ -219,7 +218,7 @@ def posPlusMinusCode(genePostig_dict, plusminus, contig_seq, plus=True, minus=Tr return geneDict -def posPlusMinus(genePostig_dict, plusminus, contig_seq, plus=True, minus=True): +def pos_plus_minus(genePostig_dict, plusminus, contig_seq, plus=True, minus=True): minimumList = [] for x in genePostig_dict: minimumList.extend(x) @@ -250,13 +249,15 @@ def posPlusMinus(genePostig_dict, plusminus, contig_seq, plus=True, minus=True): return geneDict -def ntPosNoncode(contig_seq, postig_dict, plusminus=0): +def nt_pos_noncode(contig_seq, postig_dict, plusminus=0): genes_fa_dict = {} for i, gene in enumerate(list(postig_dict.keys())): genes_fa_dict[gene] = {"sequence": "", "description": ""} if plusminus: - genes_fa_dict[gene] = posPlusMinus(postig_dict[gene], plusminus, contig_seq) + genes_fa_dict[gene] = pos_plus_minus( + postig_dict[gene], plusminus, contig_seq + ) else: minimumList = [] for x in postig_dict[gene]: @@ -269,7 +270,7 @@ def ntPosNoncode(contig_seq, postig_dict, plusminus=0): return genes_fa_dict -def negPlusMinusCode( +def neg_plus_minus_code( geneNegtig_dict, seqlen, plusminus, rev_seq, plus=True, minus=True ): minimumList, geneDict = [], {"sequence": "", "description": ""} @@ -300,7 +301,7 @@ def negPlusMinusCode( return geneDict -def negPlusMinus(neg_gene, seqlen, plusminus, rev_seq, plus=True, minus=True): +def neg_plus_minus(neg_gene, seqlen, plusminus, rev_seq, plus=True, minus=True): minimumList = [] for x in neg_gene: minimumList.extend(x) @@ -337,7 +338,7 @@ def negPlusMinus(neg_gene, seqlen, plusminus, rev_seq, plus=True, minus=True): return geneDict -def ntNeg(rev_seq, negtig_dict, plusminus=0): +def nt_neg(rev_seq, negtig_dict, plusminus=0): genes_fa_dict = {} seqlen = len(rev_seq.rstrip()) @@ -352,13 +353,13 @@ def ntNeg(rev_seq, negtig_dict, plusminus=0): return genes_fa_dict -def ntNegNoncode(rev_seq, negtig_dict, plusminus=0): +def nt_neg_noncode(rev_seq, negtig_dict, plusminus=0): genes_fa_dict = {} seqlen = len(rev_seq.rstrip()) for gene in negtig_dict: if plusminus: - genes_fa_dict[gene] = negPlusMinus( + genes_fa_dict[gene] = neg_plus_minus( negtig_dict[gene], plusminus, seqlen, rev_seq ) else: @@ -389,31 +390,33 @@ def ntmain( if fullRegion: flanks, coding = True, False contig_dict, contig_info, genes_fa_dict, geneOrder = {}, {}, {}, {} - cdss = grabCDS(sortMain(gff_dicts), spacer=spacer) + cdss = grab_cds(sort_main(gff_dicts), spacer=spacer) for cds in cdss: seqid = cds["seqid"] if seqid not in contig_dict: contig_dict[seqid] = [] geneOrder[seqid] = [] contig_dict[seqid].append(cds) - gene = re.search(gff3Comps()["Alias"], cds["attributes"])[1] + gene = re.search(gff3_comps()["Alias"], cds["attributes"])[1] if gene not in set(geneOrder[seqid]): geneOrder[seqid].append(gene) for seqid in contig_dict: contig_info[seqid] = [ ( - re.search(gff3Comps()["Alias"], contig_dict[seqid][0]["attributes"])[1], + re.search(gff3_comps()["Alias"], contig_dict[seqid][0]["attributes"])[ + 1 + ], contig_dict[seqid][0]["strand"], ), ( - re.search(gff3Comps()["Alias"], contig_dict[seqid][-1]["attributes"])[ + re.search(gff3_comps()["Alias"], contig_dict[seqid][-1]["attributes"])[ 1 ], contig_dict[seqid][-1]["strand"], ), ] - pos_dict, neg_dict = grabCoords(cdss) + pos_dict, neg_dict = grab_coords(cdss) contig_fa_dict, startFlanks, endFlanks = ( {seqid: {} for seqid in contig_info}, @@ -430,7 +433,7 @@ def ntmain( endGene, endStrand = contig_info[seqid][1][0], contig_info[seqid][1][1] if startGene == endGene: if startStrand == "+": - startFlanks[seqid] = posPlusMinusCode( + startFlanks[seqid] = pos_plus_minus_code( pos_dict[seqid][startGene], plusminus, assem_dict[seqid]["sequence"], @@ -440,7 +443,7 @@ def ntmain( rev_comp = str( Seq(assem_dict[seqid]["sequence"]).reverse_complement() ) - startFlanks[seqid] = negPlusMinusCode( + startFlanks[seqid] = neg_plus_minus_code( neg_dict[seqid][startGene], len(rev_comp), plusminus, @@ -449,7 +452,7 @@ def ntmain( del neg_dict[seqid][startGene] continue if startStrand == "+": - startFlanks[seqid] = posPlusMinusCode( + startFlanks[seqid] = pos_plus_minus_code( pos_dict[seqid][startGene], plusminus, assem_dict[seqid]["sequence"], @@ -460,7 +463,7 @@ def ntmain( rev_comp = str( Seq(assem_dict[seqid]["sequence"]).reverse_complement() ) - startFlanks[seqid] = negPlusMinusCode( + startFlanks[seqid] = neg_plus_minus_code( neg_dict[seqid][startGene], len(rev_comp), plusminus, @@ -469,7 +472,7 @@ def ntmain( ) del neg_dict[seqid][startGene] if endStrand == "+": - endFlanks[seqid] = posPlusMinusCode( + endFlanks[seqid] = pos_plus_minus_code( pos_dict[seqid][endGene], plusminus, assem_dict[seqid]["sequence"], @@ -480,7 +483,7 @@ def ntmain( rev_comp = str( Seq(assem_dict[seqid]["sequence"]).reverse_complement() ) - endFlanks[seqid] = negPlusMinusCode( + endFlanks[seqid] = neg_plus_minus_code( neg_dict[seqid][endGene], len(rev_comp), plusminus, @@ -498,7 +501,7 @@ def ntmain( if startGene == endGene: if startStrand == "+": - startFlanks[seqid] = posPlusMinus( + startFlanks[seqid] = pos_plus_minus( pos_dict[seqid][startGene], plusminus, assem_dict[seqid]["sequence"], @@ -508,7 +511,7 @@ def ntmain( rev_comp = str( Seq(assem_dict[seqid]["sequence"]).reverse_complement() ) - startFlanks[seqid] = negPlusMinus( + startFlanks[seqid] = neg_plus_minus( neg_dict[seqid][startGene], len(rev_comp), plusminus, @@ -528,18 +531,18 @@ def ntmain( else: region.append(max(neg_dict[seqid][endGene])) name = startGene + "-" + endGene - genes_fa_dict[name + "_sense"] = posPlusMinus( + genes_fa_dict[name + "_sense"] = pos_plus_minus( region, plusminus, assem_dict[seqid]["sequence"] ) rev_comp = str( Seq(assem_dict[seqid]["sequence"]).reverse_complement() ) - genes_fa_dict[name + "_antisense"] = negPlusMinus( + genes_fa_dict[name + "_antisense"] = neg_plus_minus( region, len(rev_comp), plusminus, rev_comp ) if startStrand == "+": - startFlanks[seqid] = posPlusMinus( + startFlanks[seqid] = pos_plus_minus( pos_dict[seqid][startGene], plusminus, assem_dict[seqid]["sequence"], @@ -550,7 +553,7 @@ def ntmain( rev_comp = str( Seq(assem_dict[contig]["sequence"]).reverse_complement() ) - startFlanks[seqid] = negPlusMinus( + startFlanks[seqid] = neg_plus_minus( neg_dict[seqid][startGene], len(rev_comp), plusminus, @@ -559,7 +562,7 @@ def ntmain( ) del neg_dict[seqid][startGene] if endStrand == "+": - endFlanks[seqid] = posPlusMinus( + endFlanks[seqid] = pos_plus_minus( pos_dict[seqid][endGene], plusminus, assem_dict[seqid]["sequence"], @@ -570,7 +573,7 @@ def ntmain( rev_comp = str( Seq(assem_dict[seqid]["sequence"]).reverse_complement() ) - endFlanks[seqid] = negPlusMinus( + endFlanks[seqid] = neg_plus_minus( neg_dict[seqid][endGene], len(rev_comp), plusminus, @@ -583,7 +586,7 @@ def ntmain( for contig in contig_fa_dict: if contig in pos_dict: contig_fa_dict[contig] = { - **ntPos(assem_dict[contig]["sequence"], pos_dict[contig]), + **nt_pos(assem_dict[contig]["sequence"], pos_dict[contig]), **contig_fa_dict[contig], } if contig in neg_dict: @@ -591,14 +594,14 @@ def ntmain( Seq(assem_dict[contig]["sequence"]).reverse_complement() ) contig_fa_dict[contig] = { - **ntNeg(rev_comp, neg_dict[contig]), + **nt_neg(rev_comp, neg_dict[contig]), **contig_fa_dict[contig], } elif not fullRegion: for contig in contig_fa_dict: if contig in pos_dict: contig_fa_dict[contig] = { - **ntPosNoncode( + **nt_pos_noncode( assem_dict[contig]["sequence"], pos_dict[contig] ), **contig_fa_dict[contig], @@ -608,7 +611,7 @@ def ntmain( Seq(assem_dict[contig]["sequence"]).reverse_complement() ) contig_fa_dict[contig] = { - **ntNegNoncode(rev_comp, neg_dict[contig]), + **nt_neg_noncode(rev_comp, neg_dict[contig]), **contig_fa_dict[contig], } if not fullRegion: @@ -625,18 +628,18 @@ def ntmain( def aamain(gff_dicts, assem_dict, spacer="\t"): - cdss = grabCDS(gff_dicts, spacer) - pos_dict, neg_dict = grabCoords(cdss) + cdss = grab_cds(gff_dicts, spacer) + pos_dict, neg_dict = grab_coords(cdss) genes_fa_dict = {} for contig in pos_dict: genes_fa_dict = { - **translatePos(assem_dict[contig]["sequence"], pos_dict[contig]), + **translate_pos(assem_dict[contig]["sequence"], pos_dict[contig]), **genes_fa_dict, } for contig in neg_dict: rev_comp = str(Seq(assem_dict[contig]["sequence"]).reverse_complement()) - genes_fa_dict = {**translateNeg(rev_comp, neg_dict[contig]), **genes_fa_dict} + genes_fa_dict = {**translate_neg(rev_comp, neg_dict[contig]), **genes_fa_dict} return genes_fa_dict @@ -646,7 +649,7 @@ def cli(): parser = argparse.ArgumentParser( description="Inputs MycoDB compatible gff3, assembly (optional), " + "and outputs nucleotides/proteins. Use an abstracted gene gff " - + "(acc2gff.py) to only output a smaller set of gene(s)." + + "(`mtdb accession gff`) to only output a smaller set of gene(s)." ) parser.add_argument("-g", "--gff", help='"-" for stdin', required=True) parser.add_argument("-n", "--nucleotide", action="store_true") @@ -659,6 +662,7 @@ def cli(): "-af", "--all_flanks", action="store_true", help="-n and -nc only" ) args = parser.parse_args() + setup_logging(verbose=getattr(args, "verbose", False)) if args.gff == "-": data = stdin2str() @@ -669,20 +673,18 @@ def cli(): assembly_dicts = {"input": fa2dict(format_path(args.assembly))} gff_dicts = {"input": input_gff} else: - db = mtdb(primaryDB()).set_index("ome") + db = mtdb(primary_db()).set_index("ome") gff_dicts, assembly_dicts = {}, {} try: for line in input_gff: - gene = re.search(gff3Comps()["Alias"], line["attributes"])[1] + gene = re.search(gff3_comps()["Alias"], line["attributes"])[1] ome = re.search(r"(.*?)_", gene)[1] if ome not in gff_dicts: gff_dicts[ome] = [] assembly_dicts[ome] = fa2dict(db[ome]["fna"]) gff_dicts[ome].append(line) except IndexError: - eprint( - "\nERROR: " + args.gff + " is incompatible with MycotoolsDB", flush=True - ) + logger.error("" + args.gff + " is incompatible with MycotoolsDB") sys.exit(1) if args.protein: diff --git a/mycotools/fa2mass.py b/mycotools/seq/mass.py similarity index 78% rename from mycotools/fa2mass.py rename to mycotools/seq/mass.py index ad096ad..900f45d 100755 --- a/mycotools/fa2mass.py +++ b/mycotools/seq/mass.py @@ -1,16 +1,19 @@ #! /usr/bin/env python3 -import os -import re import sys -from mycotools.lib.kontools import sys_start, format_path, fmt_float +import logging +from mycotools.lib.kontools import sys_start, format_path, fmt_float, setup_logging from mycotools.lib.biotools import fa2dict, calc_weight +logger = logging.getLogger(__name__) + + def cli(): + setup_logging() usage = "USAGE: Inputs amino acid fasta outputs linear protein weights" - args = sys_start(sys.argv, usage, 1) + sys_start(sys.argv, usage, 1) fa = fa2dict(format_path(sys.argv[1])) print("#protein\tkDa", flush=True) diff --git a/mycotools/fna2faa.py b/mycotools/seq/translate.py similarity index 89% rename from mycotools/fna2faa.py rename to mycotools/seq/translate.py index 8670141..126e98b 100755 --- a/mycotools/fna2faa.py +++ b/mycotools/seq/translate.py @@ -1,12 +1,17 @@ #! /usr/bin/env python3 import sys +import logging from Bio.Seq import Seq -from mycotools.lib.kontools import format_path, stdin2str, sys_start +from mycotools.lib.kontools import format_path, stdin2str, sys_start, setup_logging from mycotools.lib.biotools import fa2dict, dict2fa +logger = logging.getLogger(__name__) + + def cli(): + setup_logging() usage = 'Input nucleotide fasta ("-" for stdin), translate to protein fasta' args = sys_start(sys.argv[1:], usage, 1) if args[0] == "-": diff --git a/mycotools/stats/__init__.py b/mycotools/stats/__init__.py new file mode 100644 index 0000000..adcc5f1 --- /dev/null +++ b/mycotools/stats/__init__.py @@ -0,0 +1,36 @@ +#! /usr/bin/env python3 +"""Dispatcher for the `mycotools stats` subcommand. + +Routes `mycotools stats ...` to a statistics module.""" +from mycotools.lib.subcmd import Dispatcher + +# subcommand name/alias -> submodule within this package (mycotools.stats.) +SUBCOMMANDS = { + "annotation": "annotation", + "assembly": "assembly", +} + +DESCRIPTION = """Summary statistics for annotations and assemblies + +Kinds (all following arguments are forwarded to the kind): + annotation gene annotation statistics (gff3/gtf or MTDB) + assembly genome assembly statistics + +Examples: + mycotools stats annotation -h + mycotools stats assembly -h""" + +_dispatcher = Dispatcher( + "mycotools stats", + "mycotools.stats", + SUBCOMMANDS, + DESCRIPTION, + metavar="KIND", + arg_help="statistic to compute (see below)", +) +main = _dispatcher.main +cli = _dispatcher.cli + + +if __name__ == "__main__": + cli() diff --git a/mycotools/stats/__main__.py b/mycotools/stats/__main__.py new file mode 100644 index 0000000..e50908d --- /dev/null +++ b/mycotools/stats/__main__.py @@ -0,0 +1,6 @@ +#! /usr/bin/env python3 +"""Enable ``python -m mycotools.stats`` to run the stats dispatcher.""" +from mycotools.stats import cli + +if __name__ == "__main__": + cli() diff --git a/mycotools/annotationStats.py b/mycotools/stats/annotation.py similarity index 96% rename from mycotools/annotationStats.py rename to mycotools/stats/annotation.py index bc7b246..fb1c7a4 100755 --- a/mycotools/annotationStats.py +++ b/mycotools/stats/annotation.py @@ -5,10 +5,15 @@ import os import re import sys +import logging from itertools import chain from collections import defaultdict -from mycotools.lib.biotools import gff2list, gff3Comps -from mycotools.lib.kontools import format_path, eprint +from mycotools.lib.biotools import gff2list, gff3_comps +from mycotools.lib.kontools import format_path, setup_logging +from pathlib import Path + + +logger = logging.getLogger(__name__) def compile_alia(gff_path, output, ome=None): @@ -37,7 +42,7 @@ def compile_alia(gff_path, output, ome=None): for entry in gff: try: # extract the alias from the attributes field - alias = re.search(gff3Comps()["Alias"], entry["attributes"])[1] + alias = re.search(gff3_comps()["Alias"], entry["attributes"])[1] except TypeError: raise TypeError(f"entry without MTDB alias: {entry}") if entry["type"].lower() == "gene": @@ -58,7 +63,6 @@ def compile_alia(gff_path, output, ome=None): # calculate the lengths of each specific type gene_lens = sorted([v[1] - v[0] for v in chain(*list(gene_dict.values()))]) prot_lens = sorted([v[1] - v[0] for v in chain(*list(prot_dict.values()))]) - exon_lens = sorted([v[1] - v[0] for v in chain(*list(exon_dict.values()))]) mrna_lens = sorted([v[1] - v[0] for v in chain(*list(mrna_dict.values()))]) trna_lens = sorted([v[1] - v[0] for v in chain(*list(trna_dict.values()))]) orna_lens = sorted([v[1] - v[0] for v in chain(*list(orna_dict.values()))]) @@ -69,7 +73,6 @@ def compile_alia(gff_path, output, ome=None): # calculate the number of each specific type gene_len = len(gene_lens) prot_len = len(prot_lens) - exon_len = len(exon_lens) mrna_len = len(mrna_lens) trna_len = len(trna_lens) orna_len = len(orna_lens) @@ -250,7 +253,7 @@ def main(in_path, log_path=None, cpus=1, db=None): db = mtdb(in_path).set_index() prevOmes = {} - if log_path and os.path.isfile(log_path): + if log_path and Path(log_path).is_file(): with open(log_path, "r") as raw: for line in raw: if not line.startswith("#"): @@ -299,7 +302,7 @@ def main(in_path, log_path=None, cpus=1, db=None): def cli(): - output = False + setup_logging() usage = "\nUSAGE: `gff`/`gtf`/`gff3` OR mycotoolsDB, optional output file\n" if "-h " in sys.argv or "--help" in sys.argv or "-h" == sys.argv[-1]: print(usage, flush=True) @@ -307,7 +310,7 @@ def cli(): elif len(sys.argv) < 2: print(usage, flush=True) sys.exit(1) - elif not os.path.isfile(format_path(sys.argv[1])): + elif not Path(format_path(sys.argv[1])).is_file(): print(usage, flush=True) sys.exit(1) elif len(sys.argv) > 2: diff --git a/mycotools/assemblyStats.py b/mycotools/stats/assembly.py similarity index 92% rename from mycotools/assemblyStats.py rename to mycotools/stats/assembly.py index 90c0393..28bbc28 100755 --- a/mycotools/assemblyStats.py +++ b/mycotools/stats/assembly.py @@ -9,14 +9,18 @@ import os import sys -import copy +import logging import multiprocessing as mp from mycotools.lib.dbtools import mtdb from mycotools.lib.biotools import fa2dict -from mycotools.lib.kontools import format_path, eprint +from mycotools.lib.kontools import format_path, setup_logging +from pathlib import Path -def calcMask(contig_list): +logger = logging.getLogger(__name__) + + +def calc_mask(contig_list): """Calculate the percent of the sequence that is masked (lower-cased)""" seq = "".join([x["sequence"] for x in contig_list]) mask = seq.count("a") @@ -27,7 +31,7 @@ def calcMask(contig_list): return mask -def sortContigs(assembly_path): +def sort_contigs(assembly_path): """Imports fasta, creates a list of dicts for each contig length and its name. Sorts the list in descending order by length""" @@ -112,8 +116,8 @@ def n50l50(sortedContigs): if "n50-1000bp" not in out: out["n50-1000bp"] = "na" out["l50-1000bp"] = "na" - maskCount = calcMask(sortedContigs) - maskCount1000 = calcMask(pass_fa) + maskCount = calc_mask(sortedContigs) + maskCount1000 = calc_mask(pass_fa) out["mask%"] = maskCount / int(total) * 100 out["mask%-1000bp"] = maskCount1000 / int(total1000) * 100 except KeyError: # no stats acquired ?? @@ -123,7 +127,7 @@ def n50l50(sortedContigs): def mngr(assembly_path, ome): - sortedContigs = sortContigs(assembly_path) + sortedContigs = sort_contigs(assembly_path) calcs = n50l50(sortedContigs) return ome, tuple([(x, calcs[x]) for x in calcs]) @@ -141,7 +145,7 @@ def main(in_path, log_path=None, cpus=1, db=None): # parse the output file if it currently exists to avoid redundant runs prev_omes = {} if log_path: - if not os.path.isfile(log_path): + if not Path(log_path).is_file(): with open(log_path, "w") as log_open: log_open.write(head) else: @@ -176,7 +180,7 @@ def main(in_path, log_path=None, cpus=1, db=None): if res[1]: calcs[res[0]] = "\t".join([str(x[1]) for x in res[1]]) else: - eprint("\t\tERROR:\t" + ome, flush=True) + logger.error(ome) # sort the results by the ome code alphabetically calcs = { @@ -199,12 +203,12 @@ def main(in_path, log_path=None, cpus=1, db=None): # if there is not a database then run for the input file else: - sortedContigs = sortContigs(in_path) + sortedContigs = sort_contigs(in_path) calculations = n50l50(sortedContigs) if calculations: - stats[os.path.basename(os.path.abspath(in_path))] = n50l50(sortedContigs) + stats[Path(os.path.abspath(in_path)).name] = n50l50(sortedContigs) else: - eprint("\tERROR:\t" + in_path, flush=True) + logger.error(in_path) # print the stats to standard out, depending on if there are contigs # less than 1000 bp @@ -223,6 +227,7 @@ def main(in_path, log_path=None, cpus=1, db=None): def cli(): + setup_logging() usage = "\nUSAGE: assembly statistics\nAssembly `fasta` or mycotoolsDB, optional output file if using database\n" if {"-h", "--help"}.intersection(set(sys.argv)): print(usage, flush=True) diff --git a/mycotools/utils/curGFF3.py b/mycotools/utils/cur_gff3.py similarity index 95% rename from mycotools/utils/curGFF3.py rename to mycotools/utils/cur_gff3.py index eee6a52..28a514d 100755 --- a/mycotools/utils/curGFF3.py +++ b/mycotools/utils/cur_gff3.py @@ -16,14 +16,16 @@ # RNAs with no parent # CDS/exons with no parent -import os +import logging import re import sys import copy from collections import defaultdict from itertools import combinations, chain -from mycotools.lib.kontools import format_path, sys_start, eprint -from mycotools.lib.biotools import gff2list, list2gff, gff3Comps +from mycotools.lib.kontools import format_path, sys_start +from mycotools.lib.biotools import gff2list, list2gff, gff3_comps + +logger = logging.getLogger(__name__) class RNAError(Exception): @@ -53,7 +55,7 @@ def add_missing(gff_list, intron, comps, ome): "5_prime_utr", "3_prime_utr", } - out_genes, t_list, rnas, introns = {}, [], {}, {} + out_genes, rnas, introns = {}, {}, {} mtdb_count, pseudocount, alt_alias = 1, 1, {} cds2par = {} rna_changes = {} # a dictionary for changing ambigious rna id names for @@ -347,7 +349,6 @@ def add_missing(gff_list, intron, comps, ome): continue if introns: - exons = {} for gene, intronList in introns.items(): if out_genes[gene]["exon"]: continue @@ -411,7 +412,7 @@ def add_missing(gff_list, intron, comps, ome): for cds, cds_e in geneInfo["cds"].items(): for cds_d in cds_e: coords = sorted([cds_d["start"], cds_d["end"]]) - par = re.search(gff3Comps()["par"], cds_d["attributes"])[1] + par = re.search(gff3_comps()["par"], cds_d["attributes"])[1] cds_info[cds].add( ( tuple(coords), @@ -428,7 +429,6 @@ def add_missing(gff_list, intron, comps, ome): out_list = [] for geneID, geneInfo in out_genes.items(): - multiRNA = False if ( not any(x["type"] in {"RNA", "mRNA"} for x in geneInfo["rna"]) and not geneInfo["tmrna"] @@ -438,7 +438,7 @@ def add_missing(gff_list, intron, comps, ome): # but I don't like having to do most all of this because the files # are so inconsistently formatted geneInfo["tmrna"] = copy.deepcopy(geneInfo["gene"]) - id_ = re.search(gff3Comps()["id"], geneInfo["tmrna"][0]["attributes"])[1] + id_ = re.search(gff3_comps()["id"], geneInfo["tmrna"][0]["attributes"])[1] new_id = "mrna" + id_[4:] geneInfo["tmrna"][0][ "attributes" @@ -446,15 +446,13 @@ def add_missing(gff_list, intron, comps, ome): geneInfo["tmrna"][0]["type"] = "mRNA" for cds in geneInfo["cds"]: cds["attributes"] = re.sub( - gff3Comps()["par"], "Parent=" + new_id, cds["attributes"] + gff3_comps()["par"], "Parent=" + new_id, cds["attributes"] ) if geneInfo["rna"]: # if geneInfo['rna'][0]['type'] != 'mRNA' and not geneInfo['cds']: # del geneInfo['tmrna'] - if len(geneInfo["rna"]) > 1: - multiRNA = True if any(x["type"] == "mRNA" for x in geneInfo["rna"]): del geneInfo["tmrna"] # post-translational/transcriptional modification and not multiple @@ -546,7 +544,7 @@ def add_missing(gff_list, intron, comps, ome): # del geneInfo['texon'] if geneInfo["pseudo"]: for cds in geneInfo["cds"]: - cds["attributes"] = re.sub(gff3Comps()["Alias"], "", cds["attributes"]) + cds["attributes"] = re.sub(gff3_comps()["Alias"], "", cds["attributes"]) for seqType, seqEntry in geneInfo.items(): if seqType != "pseudo": @@ -569,7 +567,7 @@ def add_missing(gff_list, intron, comps, ome): def acquire_format(gff_list): - prot_comp = re.compile(gff3Comps()["id"]) + prot_comp = re.compile(gff3_comps()["id"]) gene = False for line in gff_list: if line["type"] == "gene": @@ -586,7 +584,7 @@ def acquire_format(gff_list): return None -def compile_genes(cur_list, ome, pseudocount=0, comps=gff3Comps(), cur_seqids=False): +def compile_genes(cur_list, ome, pseudocount=0, comps=gff3_comps(), cur_seqids=False): genes, pseudogenes, rnas = [], [], {} mrnas = defaultdict(list) @@ -855,7 +853,7 @@ def rename_and_organize(gff_list): seqid = entry["seqid"] if entry["type"] in {"pseudogene", "gene"}: try: - alias = re.search(gff3Comps()["Alias"], entry["attributes"])[1] + alias = re.search(gff3_comps()["Alias"], entry["attributes"])[1] except TypeError: raise TypeError(entry) alias_list = alias.split("|") @@ -874,7 +872,7 @@ def rename_and_organize(gff_list): alias2geneid[a] = gene_id geneid2alias[gene_id] = a entry["attributes"] = re.sub( - gff3Comps()["id"], "ID=" + gene_id, entry["attributes"] + gff3_comps()["id"], "ID=" + gene_id, entry["attributes"] ) scaf2gene2entries[seqid][gene_id]["gene"] = [entry] todel.append(i) @@ -884,14 +882,14 @@ def rename_and_organize(gff_list): "three_prime_utr", "five_prime_utr", }: - par = re.search(gff3Comps()["par"], entry["attributes"]) + par = re.search(gff3_comps()["par"], entry["attributes"]) if par is None: - alias = re.search(gff3Comps()["Alias"], entry["attributes"])[1] + alias = re.search(gff3_comps()["Alias"], entry["attributes"])[1] new_id = entry["type"].lower() + "_" + alias alias2geneid[alias] = new_id geneid2alias[new_id] = alias entry["attributes"] = re.sub( - gff3Comps()["id"], f"ID={new_id}", entry["attributes"] + gff3_comps()["id"], f"ID={new_id}", entry["attributes"] ) scaf2gene2entries[seqid][new_id]["gene"] = [entry] todel.append(i) @@ -904,17 +902,17 @@ def rename_and_organize(gff_list): for i, entry in enumerate(gff_list): if "RNA" in entry["type"]: - alias = re.search(gff3Comps()["Alias"], entry["attributes"])[1] + alias = re.search(gff3_comps()["Alias"], entry["attributes"])[1] rna_id = entry["type"].lower() + "_" + alias entry["attributes"] = re.sub( - gff3Comps()["id"], "ID=" + rna_id, entry["attributes"] + gff3_comps()["id"], "ID=" + rna_id, entry["attributes"] ) try: gene_id = alias2geneid[alias] except KeyError: raise KeyError("RNA alias that does not link with gene") entry["attributes"] = re.sub( - gff3Comps()["par"], "Parent=" + gene_id, entry["attributes"] + gff3_comps()["par"], "Parent=" + gene_id, entry["attributes"] ) seqid = entry["seqid"] scaf2gene2entries[seqid][gene_id]["rna"][alias] = [entry] @@ -928,7 +926,7 @@ def rename_and_organize(gff_list): id_dict = defaultdict(dict) for entry in gff_list: try: - alias = re.search(gff3Comps()["Alias"], entry["attributes"])[1] + alias = re.search(gff3_comps()["Alias"], entry["attributes"])[1] except TypeError: raise TypeError(entry) typ = entry["type"].lower() @@ -938,23 +936,23 @@ def rename_and_organize(gff_list): id_dict[alias][typ] += 1 oth_id = typ + str(id_dict[alias][typ]) + "_" + alias entry["attributes"] = re.sub( - gff3Comps()["id"], "ID=" + oth_id, entry["attributes"] + gff3_comps()["id"], "ID=" + oth_id, entry["attributes"] ) if alias in alias2rnaid: par_id = alias2rnaid[alias] gene_id = alias2geneid[alias] entry["attributes"] = re.sub( - gff3Comps()["par"], "Parent=" + par_id, entry["attributes"] + gff3_comps()["par"], "Parent=" + par_id, entry["attributes"] ) scaf2gene2entries[entry["seqid"]][gene_id]["rna"][alias].append(entry) elif alias in alias2geneid: par_id = alias2geneid[alias] entry["attributes"] = re.sub( - gff3Comps()["par"], "Parent=" + par_id, entry["attributes"] + gff3_comps()["par"], "Parent=" + par_id, entry["attributes"] ) scaf2gene2entries[entry["seqid"]][par_id]["gene"].append(entry) else: - par_id = re.search(gff3Comps()["par"], entry["attributes"])[1] + par_id = re.search(gff3_comps()["par"], entry["attributes"])[1] if par_id in rnaid2alias: new_alias = rnaid2alias[par_id] gene_id = alias2geneid[new_alias] @@ -964,10 +962,10 @@ def rename_and_organize(gff_list): else: raise KeyError(f"missing parent: {entry}") entry["attributes"] = re.sub( - gff3Comps()["par"], "Parent=" + par_id, entry["attributes"] + gff3_comps()["par"], "Parent=" + par_id, entry["attributes"] ) entry["attributes"] = re.sub( - gff3Comps()["Alias"], "Alias=" + new_alias, entry["attributes"] + gff3_comps()["Alias"], "Alias=" + new_alias, entry["attributes"] ) scaf2gene2entries[entry["seqid"]][gene_id]["etc"][new_alias].append(entry) @@ -995,14 +993,14 @@ def rename_and_organize(gff_list): return out_gff -def curGff3(gff_list, ome, cur_seqids=False): +def cur_gff3(gff_list, ome, cur_seqids=False): cur_list, intron = [], False for line in gff_list: if line["type"] == "intron": intron = True - cur_list, pseudocount = add_missing(gff_list, intron, gff3Comps(), ome) + cur_list, pseudocount = add_missing(gff_list, intron, gff3_comps(), ome) final_list = compile_genes(cur_list, ome, pseudocount, cur_seqids=cur_seqids) return final_list @@ -1018,10 +1016,10 @@ def main(gff_path, ome, cur_seqids=False): typ = True if not typ: - eprint("\tERROR: type unknown ", flush=True) + logger.error("type unknown ") return None - new_gff = curGff3(gff, ome, cur_seqids) + new_gff = cur_gff3(gff, ome, cur_seqids) clean_gff = rename_and_organize(new_gff) return clean_gff diff --git a/mycotools/utils/extractHmmAcc.py b/mycotools/utils/extract_hmm_acc.py similarity index 73% rename from mycotools/utils/extractHmmAcc.py rename to mycotools/utils/extract_hmm_acc.py index dcc0bfd..3aff23b 100755 --- a/mycotools/utils/extractHmmAcc.py +++ b/mycotools/utils/extract_hmm_acc.py @@ -1,12 +1,15 @@ #! /usr/bin/env python3 -import os +import logging import re import sys -from mycotools.lib.kontools import eprint, file2list +from mycotools.lib.kontools import file2list +from pathlib import Path +logger = logging.getLogger(__name__) -def grabAccs(db_str): + +def grab_accs(db_str): accessions = [] hmm_search = re.compile(r"HMMER\d\/f \[.*?\]\nNAME(.*?)\nACC +(.*?)\n[^\/]*?\/\/") @@ -23,7 +26,7 @@ def grabAccs(db_str): return accessions -def hmmExtract(accession, db_str): +def hmm_extract(accession, db_str): hmm_search = re.search( r"HMMER\d\/f \[.*?\].*?\n^NAME.*?\n^ACC " + accession + r"[\s\S]*?^\/\/", @@ -35,10 +38,7 @@ def hmmExtract(accession, db_str): r"HMMER\d\/f \[.*?\]\nNAME +" + accession + r"[\s\S]*?\/\/", db_str ) if not hmm_search: - eprint( - "\nERROR: " + accession + " does not exist or unexpected error\n", - flush=True, - ) + logger.error("" + accession + " does not exist or unexpected error\n") sys.exit(2) hmm = hmm_search[0] + "\n" @@ -49,23 +49,23 @@ def hmmExtract(accession, db_str): def main(hmm_db, accessions=False): if accessions: - if os.path.isfile(accessions): + if Path(accessions).is_file(): accessions = file2list(accessions) hmm_str = "" for accession in accessions: - hmm_str += hmmExtract(accession, hmm_db) + hmm_str += hmm_extract(accession, hmm_db) else: - hmm_str = hmmExtract(accessions, hmm_db) + hmm_str = hmm_extract(accessions, hmm_db) else: - accessions = grabAccs(hmm_db) + accessions = grab_accs(hmm_db) hmm_str = {} if len(accessions[0]) == 3: accessions = [x[2] for x in accessions] elif len(accessions[0]) == 2: accessions = [x[1] for x in accessions] for accession in accessions: - hmm_str[accession] = hmmExtract(accession, hmm_db) + hmm_str[accession] = hmm_extract(accession, hmm_db) return hmm_str @@ -80,22 +80,22 @@ def cli(): print(usage, flush=True) sys.exit(1) - if not os.path.isfile(sys.argv[1]): - eprint("\nERROR: Invalid `.hmm` database path", flush=True) + if not Path(sys.argv[1]).is_file(): + logger.error("Invalid `.hmm` database path") sys.exit(2) - print("\nReading hmm database ...", flush=True) + logger.info("Reading hmm database ...") with open(sys.argv[1], "r") as raw_hmm_db: hmm_db = raw_hmm_db.read() if len(sys.argv) > 2: hmm_str = main(hmm_db, accessions=sys.argv[2]) - print("\nWriting abstracted hmms ...", flush=True) + logger.info("Writing abstracted hmms ...") with open(sys.argv[2] + ".hmm", "w") as out: out.write(hmm_str) else: - if not os.path.isdir("hmm"): - os.mkdir("hmm") + if not Path("hmm").is_dir(): + Path("hmm").mkdir() hmm_strs = main(hmm_db) for accession in hmm_strs: with open("hmm/" + accession + ".hmm", "w") as out: diff --git a/mycotools/utils/extractHmmsearch.py b/mycotools/utils/extract_hmmsearch.py similarity index 93% rename from mycotools/utils/extractHmmsearch.py rename to mycotools/utils/extract_hmmsearch.py index 7cc6589..f6e3f91 100755 --- a/mycotools/utils/extractHmmsearch.py +++ b/mycotools/utils/extract_hmmsearch.py @@ -1,10 +1,20 @@ #! /usr/bin/env python3 -import os +import logging import re import sys import argparse -from mycotools.lib.kontools import intro, outro, file2list, format_path, mkOutput +from mycotools.lib.kontools import ( + intro, + outro, + file2list, + format_path, + mk_output, + setup_logging, +) +from pathlib import Path + +logger = logging.getLogger(__name__) def grab_names(data, query=False): @@ -98,10 +108,7 @@ def grab_hits( if x != "" and x != ".." and "[" not in x and "]" not in x ] if len(outLine) != 13: - print( - "\nINVALID ALIGNMENT HEADERS - CHECK SCRIPT/ALIGNMENT\n", - flush=True, - ) + logger.error("INVALID ALIGNMENT HEADERS - CHECK SCRIPT/ALIGNMENT") sys.exit(5) if threshold: hmmStart = int(outLine[6]) @@ -196,12 +203,12 @@ def grab_hits( else: hit_str, aln_str = None, None - print("did not work", flush=True) + logger.error("did not work") return hit_str, aln_str -def synthesizeHits(out_dict): +def synthesize_hits(out_dict): check = [] for hit in out_dict: @@ -268,7 +275,7 @@ def main(data, accession, best, threshold, evalue, bitscore, query=True, header= align += t_aligns out[x] = (hits, align) else: - htis, align = "", "" + align = "" if header: hits = ( "#seq\tseq_e\tseq_score\tseq_bias\tdom_e\tdom_score\tdom_bias\texp\tN\n" @@ -327,13 +334,14 @@ def cli(): "-o", "--output", help="Output file name/path (extensions automatically applied" ) args = parser.parse_args() + setup_logging(verbose=getattr(args, "verbose", False)) # initialize output file structure output = format_path(args.output) if not args.output: - output = mkOutput(os.getcwd() + "/", "extractHmmsearch") - elif not os.path.isdir(output): - os.mkdir(args.output) + output = mk_output(str(Path.cwd()) + "/", "extract_hmmsearch") + elif not Path(output).is_dir(): + Path(args.output).mkdir() args_dict = { "Input": args.input, @@ -352,24 +360,24 @@ def cli(): args.accession = True else: if args.query: - if os.path.isfile(args.query): + if Path(args.query).is_file(): ques = file2list(args.query) else: ques = [args.query] args.query = True elif args.accession: - if os.path.isfile(args.accession): + if Path(args.accession).is_file(): ques = file2list(args.accession) else: ques = [args.accession] args.accession = True else: - print("\nNeed `-q` or `-a` specified\n", flush=True) + logger.error("Need `-q` or `-a` specified") sys.exit(8) - if not os.path.isfile(args.input): - print("\n\tNot a valid input file\n", flush=True) - sys.eixt(2) + if not Path(args.input).is_file(): + logger.error("Not a valid input file") + sys.exit(2) with open(args.input, "r") as raw: data = raw.read() @@ -398,7 +406,7 @@ def cli(): query=args.query, ) out_dict = {**out_dict, **temp_out_dict} - out_dict = synthesizeHits(out_dict) + out_dict = synthesize_hits(out_dict) for name in out_dict: with open(output + "/" + name + ".hits.tsv", "w") as out: diff --git a/mycotools/utils/gff2gff3.py b/mycotools/utils/gff2gff3.py index 6ffa950..14f1ecf 100755 --- a/mycotools/utils/gff2gff3.py +++ b/mycotools/utils/gff2gff3.py @@ -2,20 +2,22 @@ # NEED to arrive at a consensus for protein and transcript IDs +import logging import re import sys -import copy import argparse from collections import defaultdict -from mycotools.lib.biotools import gff2list, list2gff, gff2Comps, gff3Comps -from mycotools.lib.kontools import format_path, eprint, vprint -from mycotools.utils.gtf2gff3 import add_genes, remove_start_stop -from mycotools.utils.curGFF3 import rename_and_organize +from mycotools.lib.biotools import gff2list, list2gff, gff2_comps, gff3_comps +from mycotools.lib.kontools import format_path, setup_logging +from mycotools.utils.gtf2gff3 import add_genes +from mycotools.utils.cur_gff3 import rename_and_organize + +logger = logging.getLogger(__name__) def gff2gff3(gff_list, ome, jgi_ome): - comps2, exon_dict, cds_dict, out_list, gene_dict = gff2Comps(), {}, {}, [], {} + comps2, exon_dict, cds_dict, out_list, gene_dict = gff2_comps(), {}, {}, [], {} for entry in gff_list: if entry["type"] == "exon": name = re.search(comps2["id"], entry["attributes"])[1] @@ -49,7 +51,7 @@ def gff2gff3(gff_list, ome, jgi_ome): cds_id = "CDS_$_" + str(cds_dict[name]) entry["attributes"] = "ID=" + cds_id + ";Alias=" + name - comps3 = gff3Comps() + comps3 = gff3_comps() for entry in gff_list: if entry["type"] not in {"start_codon", "stop_codon"}: if entry["type"] == "exon": @@ -128,12 +130,12 @@ def resolve_alternate_splicing(gff): # this is a hack job and should be done during add_genes contig2gene, a2z, a2gi = defaultdict(dict), defaultdict(list), {} for i, entry in enumerate(gff): - alias = re.search(gff3Comps()["Alias"], entry["attributes"])[1] + alias = re.search(gff3_comps()["Alias"], entry["attributes"])[1] if entry["type"] == "gene": a2gi[alias] = i contig = entry["seqid"] start, end = entry["start"], entry["end"] - gene = re.search(gff3Comps()["id"], entry["attributes"])[1] + gene = re.search(gff3_comps()["id"], entry["attributes"])[1] contig2gene[contig][gene] = (start, end, alias) a2z[alias].append(entry) @@ -172,7 +174,7 @@ def resolve_alternate_splicing(gff): max_iz, min_iz = [], [] for g_entry in a2z[a0]: if g_entry["type"] == "gene": - gid = re.search(gff3Comps()["id"], g_entry["attributes"])[1] + gid = re.search(gff3_comps()["id"], g_entry["attributes"])[1] g_entry["attributes"] += "|" + "|".join(accs[1:]) max_iz.append(max(g_entry["start"], g_entry["end"])) min_iz.append(min(g_entry["start"], g_entry["end"])) @@ -182,7 +184,7 @@ def resolve_alternate_splicing(gff): for entry in a2z[a1]: if "RNA" in entry["type"]: entry["attributes"] = re.sub( - gff3Comps()["par"], f"Parent={gid}", entry["attributes"] + gff3_comps()["par"], f"Parent={gid}", entry["attributes"] ) max_iz.append(max(entry["start"], entry["end"])) min_iz.append(min(entry["start"], entry["end"])) @@ -197,7 +199,7 @@ def resolve_alternate_splicing(gff): def find_jgi_problems(gff3): # check for out of bounds CDS and exons rna2gene = {} - gene_coords, other_coords, comps = {}, defaultdict(list), gff3Comps() + gene_coords, other_coords, comps = {}, defaultdict(list), gff3_comps() for entry in gff3: if entry["type"] == "gene": gene = re.search(comps["id"], entry["attributes"])[1] @@ -227,23 +229,18 @@ def find_jgi_problems(gff3): def main(gff_list, ome, jgi_ome, safe=True, verbose=True): if gff_list[0]["attributes"].startswith("gene_id"): - comps = gtfComps() + comps = gtf_comps() gene_prefix = "gene_id" else: - comps = gff2Comps() + comps = gff2_comps() gene_prefix = "name" gff_prep, failed, flagged = add_genes( gff_list, safe=safe, comps=comps, gene_prefix=gene_prefix ) if failed: - vprint(str(len(failed)) + "\tgenes failed", v=verbose, e=True, flush=True) + logger.debug(str(len(failed)) + "genes failed") if flagged: - vprint( - str(len(flagged)) + "\tgene coordinates from exons", - v=verbose, - e=True, - flush=True, - ) + logger.debug(str(len(flagged)) + "gene coordinates from exons") pregff3 = gff2gff3(gff_prep, ome, jgi_ome) gff3 = resolve_alternate_splicing(pregff3) gff3 = rename_and_organize(gff3) @@ -253,11 +250,11 @@ def main(gff_list, ome, jgi_ome, safe=True, verbose=True): err_name.append(err.upper()) if verbose: if err == "ob": - eprint("ERROR: genes with out of bounds coordinates", flush=True) - eprint(",".join(err_list), flush=True) + logger.error("genes with out of bounds coordinates") + logger.info(",".join(err_list)) elif err == "nr": - eprint("ERROR: missing RNA", flush=True) - eprint(",".join(err_list), flush=True) + logger.error("missing RNA") + logger.info(",".join(err_list)) return gff3, errors @@ -275,9 +272,10 @@ def cli(): help="Fail genes w/o CDS sequences that lack start or stop codons", ) args = parser.parse_args() + setup_logging(verbose=getattr(args, "verbose", False)) gff_list = gff2list(format_path(args.input)) - eprint(args.ome + "\t" + args.input, flush=True) + logger.info(args.ome + "" + args.input) gff3, errors = main(gff_list, args.ome, args.jgi, args.fail) print(list2gff(gff3), flush=True) diff --git a/mycotools/utils/gtf2gff3.py b/mycotools/utils/gtf2gff3.py index 7e381c6..87e5995 100755 --- a/mycotools/utils/gtf2gff3.py +++ b/mycotools/utils/gtf2gff3.py @@ -4,7 +4,7 @@ # NEED to arrive at a consensus for protein IDs # NEED to name tRNAs, pseudogenes, etc. similar to curGFF.py -import os +import logging import re import sys import copy @@ -16,16 +16,19 @@ list2gff, fa2dict, dict2fa, - gtfComps, - gff3Comps, - gff2Comps, + gtf_comps, + gff3_comps, + gff2_comps, ) -from mycotools.lib.kontools import collect_files, eprint, format_path -from mycotools.gff2seq import aamain as gff2proteome -from mycotools.utils.curGFF3 import rename_and_organize +from mycotools.lib.kontools import collect_files, format_path, setup_logging +from mycotools.seq.gff import aamain as gff2proteome +from mycotools.utils.cur_gff3 import rename_and_organize +from pathlib import Path +logger = logging.getLogger(__name__) -def grabOutput(output_pref): + +def grab_output(output_pref): """ Inputs: orthofiller `output_path` for results Outputs: gff_dict and fasta_dict of results @@ -34,44 +37,29 @@ def grabOutput(output_pref): both. """ - output_path, pref = os.path.dirname(output_pref), os.path.basename(output_pref) + output_path, pref = str(Path(output_pref).parent), Path(output_pref).name output_files = collect_files(output_path, "*") - hits = [x for x in output_files if os.path.basename(x).startswith(pref)] + hits = [x for x in output_files if Path(x).name.startswith(pref)] proteome = [x for x in hits if x.endswith(".results.aa.fasta")] gtf = [x for x in hits if x.endswith(".results.gtf")] if len(gtf) != 1 or len(proteome) != 1: if len(gtf) < 1 or len(proteome) < 1: - print("\nERROR: complete output files not detected", flush=True) + logger.error("complete output files not detected") else: - print("\nERROR: multiple eligible complete output detected", flush=True) + logger.error("multiple eligible complete output detected") sys.exit(3) return gff2list(gtf[0]), fa2dict(proteome[0]) def intron2exon(gff, gene_comp=re.compile(r"gene_id \"(.*?)\"")): - """ - Inputs: gff_dict - Outputs: gff_dict with introns converted to exons - For each entry in the gff_dict, search for the gene ID. If the gene is not - in `gene_info` add the gene as a key and populate a blank intron list, - start codon list, stop codon list, strand string, and raw list. If the - entry type is within the `gene_info` dict for the gene, then append the - list of start and stop coordinates. Append the entire entry to the raw data - key. - Create a dictionary `intron_genes` for each gene in gene_info if there is - an intron entry. For each gene in `intron_genes` create a blank list for - `exon_coords` dict under the key `gene`. If the intron coordinates' start - codon end coordinate is greater than the start coordinate, then change the - start codon entry in `intron_genes[gene]` to have the greater value first. - Repeat for the stop codon. Then sort the intron coordinates of that gene. - Append the appropriate exon coordinates for the intron based upon strand - sense. - For each gene in the gff, if it is not in `exon_coords` then simply append - to the `new_gff`. Then add genes with new exons. - """ + """Convert intron features to exons in a gff_dict. + + Groups entries by gene ID, infers each gene's exon coordinates from its + introns (respecting strand sense), and returns a gff with introns replaced + by the inferred exons.""" - comps = gtfComps() + comps = gtf_comps() gff1, gene_info = [], {} for entry in gff: try: @@ -84,7 +72,7 @@ def intron2exon(gff, gene_comp=re.compile(r"gene_id \"(.*?)\"")): continue gene_comp = re.compile(r"name \"(.*?)\"") gene = gene_comp.search(entry["attributes"])[1] - comps = gff2Comps() + comps = gff2_comps() if int(entry["start"]) > int(entry["end"]): start, end = copy.deepcopy(entry["start"]), copy.deepcopy(entry["end"]) entry["start"], entry["end"] = end, start @@ -209,7 +197,7 @@ def intron2exon(gff, gene_comp=re.compile(r"gene_id \"(.*?)\"")): return gff2, comps -def curCDS(gff, gene_compile=re.compile(r"gene_id \"(.*?)\"")): +def cur_cds(gff, gene_compile=re.compile(r"gene_id \"(.*?)\"")): new_gff, info_dict = [], {} for entry in gff: @@ -264,7 +252,7 @@ def curCDS(gff, gene_compile=re.compile(r"gene_id \"(.*?)\"")): return new_gff -def liberalRemoval(gene_dict_prep, contigs): +def liberal_removal(gene_dict_prep, contigs): check_contigs, failed, gene_dict = {}, [], {} for i in contigs: @@ -322,7 +310,7 @@ def liberalRemoval(gene_dict_prep, contigs): return gene_dict, failed -def conservativeRemoval(gene_dict_prep): +def conservative_removal(gene_dict_prep): gene_dict, flagged, failed = {}, [], [] for gene, temp in gene_dict_prep.items(): @@ -357,7 +345,7 @@ def conservativeRemoval(gene_dict_prep): ] flagged.append(gene) except IndexError: - eprint(gene + " cannot create gene coordinates", flush=True) + logger.info(gene + " cannot create gene coordinates") continue gene_dict[gene] = temp @@ -380,9 +368,8 @@ def fill_transcripts(gene_dict_prep): return gene_dict_prep -def add_genes(gtf, safe=True, comps=gtfComps(), gene_prefix="gene_id"): +def add_genes(gtf, safe=True, comps=gtf_comps(), gene_prefix="gene_id"): - contigs = defaultdict(dict) tran_compile = re.compile(comps["transcript"]) gene_compile = re.compile(comps["id"]) gene_dict_prep, gene_dict = {}, {} @@ -421,7 +408,7 @@ def add_genes(gtf, safe=True, comps=gtfComps(), gene_prefix="gene_id"): gene_dict_prep[gene]["rna"][tran] = gtf[i] gene_dict_prep = fill_transcripts(gene_dict_prep) - gene_dict, flagged, failed = conservativeRemoval(gene_dict_prep) + gene_dict, flagged, failed = conservative_removal(gene_dict_prep) check, insert_list = set(), [] for index, entry in enumerate(gtf): @@ -508,7 +495,7 @@ def remove_start_stop(gtf): return [x for x in gtf if x["type"] not in {"start_codon", "stop_codon"}] -def curate(gff, prefix=None, failed=set(), comps=gtfComps()): +def curate(gff, prefix=None, failed=set(), comps=gtf_comps()): failed = set(failed) gene_comp = re.compile(comps["id"]) @@ -526,7 +513,7 @@ def curate(gff, prefix=None, failed=set(), comps=gtfComps()): try: gene_id = gene_comp.search(line["attributes"])[1] except TypeError: - gene_id = re.search(gtfComps()["id"], line["attributes"])[1] + gene_id = re.search(gtf_comps()["id"], line["attributes"])[1] try: tran_id = tran_comp.search(line["attributes"])[1] except TypeError: @@ -561,7 +548,7 @@ def curate(gff, prefix=None, failed=set(), comps=gtfComps()): alias_dict[trans] = prefix + "_" + str(count) count += 1 - crudesortGff = preSortGFF(gff, gene_comp) + crudesortGff = pre_sort_gff(gff, gene_comp) newGff, exon_check, cds_check = [], defaultdict(int), defaultdict(int) for entry in crudesortGff: @@ -587,7 +574,6 @@ def curate(gff, prefix=None, failed=set(), comps=gtfComps()): entry["attributes"] += f";Alias={alias_dict[trans]}" exon_check[trans] += 1 elif entry["type"] == "CDS": - alias = alias_dict[trans] trans = tran_comp.search(entry["attributes"])[1] cds = cds_check[trans] + 1 entry["attributes"] = f"ID={trans}.cds{cds};Parent=" @@ -611,7 +597,7 @@ def curate(gff, prefix=None, failed=set(), comps=gtfComps()): return newGff, translation_str -def sortGene(sorting_group): +def sort_gene(sorting_group): out_group = [] for entryType in ["gene", "mrna", "trna", "rrna", "exon", "cds"]: @@ -629,7 +615,7 @@ def sortGene(sorting_group): return out_group -def sortContig(contigData): +def sort_contig(contigData): coordinates = {} for gene in contigData: @@ -647,9 +633,9 @@ def sortContig(contigData): return outContig -def preSortGFF(unsorted_gff, idComp): +def pre_sort_gff(unsorted_gff, idComp): - sorting_groups, oldGene = {}, None + sorting_groups = {} for i, entry in enumerate(unsorted_gff): seqid = entry["seqid"] if seqid not in sorting_groups: @@ -664,15 +650,15 @@ def preSortGFF(unsorted_gff, idComp): for seqid in sorting_groups: contigData = {} for gene in sorting_groups[seqid]: - contigData[gene] = sortGene(sorting_groups[seqid][gene]) - sortedGff.extend(sortContig(contigData)) + contigData[gene] = sort_gene(sorting_groups[seqid][gene]) + sortedGff.extend(sort_contig(contigData)) return sortedGff -def sortGFF(unsorted_gff, idComp): +def sort_gff(unsorted_gff, idComp): - sorting_groups, oldGene = {}, None + sorting_groups = {} for i, entry in enumerate(unsorted_gff): if entry["type"].lower() not in { "mrna", @@ -700,13 +686,13 @@ def sortGFF(unsorted_gff, idComp): for seqid in sorting_groups: contigData = {} for gene in sorting_groups[seqid]: - contigData[gene] = sortGene(sorting_groups[seqid][gene]) - sortedGff.extend(sortContig(contigData)) + contigData[gene] = sort_gene(sorting_groups[seqid][gene]) + sortedGff.extend(sort_contig(contigData)) return sortedGff -def addExons(gff): +def add_exons(gff): exon_check = {} for i in range(len(gff)): @@ -738,15 +724,15 @@ def main(gff_path, prefix, fail=True): gff = gff_path exonGtf, comps = intron2exon(gff) - exonGtfCur = curCDS(exonGtf, re.compile(comps["id"])) + exonGtfCur = cur_cds(exonGtf, re.compile(comps["id"])) exonGtfCurGenes, failed, flagged = add_genes(exonGtfCur, safe=fail, comps=comps) preGff = remove_start_stop(exonGtfCurGenes) unsortedGff, trans_str = curate(preGff, prefix, failed, comps) - # unsortedGff = addExons(gffUncur) + # unsortedGff = add_exons(gffUncur) if prefix: - out_gff = sortGFF(unsortedGff, re.compile(gff3Comps()["Alias"])) + out_gff = sort_gff(unsortedGff, re.compile(gff3_comps()["Alias"])) for i, entry in enumerate(out_gff): entry["start"] = int(entry["start"]) entry["end"] = int(entry["end"]) @@ -757,14 +743,14 @@ def main(gff_path, prefix, fail=True): return gff, trans_str, failed, flagged -def sortMain(gff, prefix): +def sort_main(gff, prefix): - id_comp = re.compile(gff3Comps()["id"]) + id_comp = re.compile(gff3_comps()["id"]) crude_sort = sorted( gff, key=lambda x: int(re.search(r"ID=" + prefix + r"_(\d+)", x["attributes"])[1]), ) - gff = preSortGFF(crude_sort, id_comp) + gff = pre_sort_gff(crude_sort, id_comp) return gff @@ -789,18 +775,19 @@ def cli(): # parser.add_argument('--fail', default = True, action = 'store_false', # help = 'Fail genes without start or stop codons.') args = parser.parse_args() + setup_logging(verbose=getattr(args, "verbose", False)) if args.output: output = format_path(args.output) - if not os.path.isdir(output): - os.mkdir(output) + if not Path(output).is_dir(): + Path(output).mkdir() output += "/" output += args.prefix else: output = args.prefix if args.sort: - gff = sortMain(gff2list(format_path(args.gff)), args.prefix) + gff = sort_main(gff2list(format_path(args.gff)), args.prefix) with open(output + ".gff3", "w") as out: out.write(list2gff(gff) + "\n") sys.exit(0) @@ -810,7 +797,7 @@ def cli(): if args.prefix: if "_" in args.prefix: - eprint('\nERROR: "_" not allowed in prefix\n', flush=True) + logger.error('\nERROR: "_" not allowed in prefix\n') sys.exit(1) gff, trans_str, failed, flagged = main( @@ -827,11 +814,11 @@ def cli(): with open(output + ".transitions", "w") as out: out.write(trans_str) if failed: - print("\n" + str(len(failed)) + " failures\n", flush=True) + logger.debug("" + str(len(failed)) + " failures\n") with open(output + ".failed", "w") as out: out.write("\n".join(["\t".join(x) for x in failed])) if flagged: - print("\n" + str(len(flagged)) + " flagged\n", flush=True) + logger.debug("" + str(len(flagged)) + " flagged\n") with open(output + ".flagged", "w") as out: out.write("\n".join(flagged)) diff --git a/mycotools/utils/jgi2db.py b/mycotools/utils/jgi2db.py index 1addbfe..43b9a1c 100755 --- a/mycotools/utils/jgi2db.py +++ b/mycotools/utils/jgi2db.py @@ -1,33 +1,30 @@ #! /usr/bin/env python3 +import logging import os import re import sys import copy -import time -import getpass import datetime import argparse -import subprocess import pandas as pd import numpy as np -from io import StringIO -from mycotools.predb2mtdb import main as predb2mtdb -from mycotools.lib.kontools import intro, outro, eprint -from mycotools.lib.dbtools import db2df, df2db, readLog, log_editor -from mycotools.jgiDwnld import jgi_login as jgi_login -from mycotools.jgiDwnld import retrieve_xml as retrieve_xml -from mycotools.jgiDwnld import jgi_dwnld as jgi_dwnld +from mycotools.lib.kontools import intro, outro, setup_logging +from mycotools.lib.dbtools import db2df, df2db, read_log, log_editor +from mycotools.download.jgi import main as jgi_dwnld +from pathlib import Path +logger = logging.getLogger(__name__) -def compileLog(log_path): + +def compile_log(log_path): log = {} - if not os.path.isfile(log_path): + if not Path(log_path).is_file(): with open(log_path, "w") as out: out.write("#assembly_acc\tfna\tgff3\tfaa") else: - log = readLog(log_path) + log = read_log(log_path) return log @@ -62,7 +59,7 @@ def jgi_redundancy_check(db, jgi_df, duplicates={}, ome_col="portal", jgi2ncbi={ else: try: db_version = float(db["version"][ome]) - except (ValueError, AttributeError) as e: + except (ValueError, AttributeError): db_version = float(db["version"][ome].replace("v", "")) db.at[ome, "version"] = db_version if version > db_version and db["source"][ome] == "jgi": @@ -136,12 +133,24 @@ def jgi_redundancy_check(db, jgi_df, duplicates={}, ome_col="portal", jgi2ncbi={ return jgi_df, db, updates, old_omes +def logged_file(log, ome, typ, output): + """Return the path of an ome's `typ` file recorded in a previous run's log, + if that file is still on disk; otherwise None. JGI files arrive gzipped and + curation decompresses them in place, so either form is accepted.""" + basename = log.get(ome, {}).get(typ, "na") + if basename in {"na", "error", "pending", ""}: + return None + path = f"{output}/{typ}/{basename}" + for candidate in (path, re.sub(r"\.gz$", "", path)): + if Path(candidate).is_file(): + return candidate + return None + + def runjgi_dwnld( jgi_df, - i, user, pwd, - ome_set, ome_col, output, log, @@ -151,82 +160,116 @@ def runjgi_dwnld( rerun, masked, spacer="\t\t", + restore_wait=None, ): - - ome = jgi_df[ome_col][i] - jgi_login(user, pwd) - - ran_dwnld = False - if ome not in ome_set: - print(spacer + "\t" + ome + ": " + jgi_df["name"][i], flush=True) - for typ in dwnlds: - if log[ome][typ] == "error": - if not rerun: - print(spacer + "\t\t" + typ + ": ERROR", flush=True) - continue - check, preexisting, new_typ, ran_dwnld, org_name = jgi_dwnld( - ome, typ, output, masked=masked - ) - if not isinstance(check, int): - jgi_df.at[i, new_typ + "_path"] = check - base_check = os.path.basename(os.path.abspath(check)) - print(spacer + "\t\t" + new_typ + ": " + str(base_check), flush=True) - log[ome][typ] = base_check - else: - print(spacer + "\t\t" + new_typ + ": ERROR", flush=True) - log[ome][typ] = "error" - if typ in {"gff3", "fna"}: - log_editor( - log_path, - ome, - ome - + "\t" - + log[ome]["fna"] - + "\t" - + log[ome]["gff3"] - + "\t" - + log[ome]["faa"], - ) - failed.append([ome, jgi_df["version"][i]]) - jgi_df = jgi_df.drop(i) - if ran_dwnld: - time.sleep(60) - break - if ran_dwnld: - time.sleep(60) - log_editor( - log_path, - ome, - ome - + "\t" - + log[ome]["fna"] - + "\t" - + log[ome]["gff3"] - + "\t" - + log[ome]["faa"], + """Download the MycoCosm portals in `jgi_df` via the JGI Data Portal API, + skipping omes a previous run already completed and recording each outcome in + the resume log. Returns (jgi_df, log, failed, deferred). + + Tape-archived genomes are skipped on the first pass and revisited once every + portal has been visited - nearly every MycoCosm genome needs a restore, so + waiting on each in turn would stall the update. `restore_wait` caps how many + minutes any one genome is waited on there; None waits for as long as JGI + takes. + + Omes without an assembly or gff3 are dropped from `jgi_df`. Genuine failures + (portal absent from MycoCosm, no such file type, corrupt download) are logged + as `error` and appended to `failed` as [ome, version], as the legacy per-ome + loop did. Omes whose files are merely awaiting a JGI tape restore are instead + logged as `pending` and returned in `deferred` - they are retried on the next + run rather than blacklisted.""" + + # resume: an ome needs no download when every requested file is logged and + # present - or, without `rerun`, was previously logged as unobtainable + todwnld_i, preexisting = [], {} + for i, row in jgi_df.iterrows(): + ome = row[ome_col] + paths = {typ: logged_file(log, ome, typ, output) for typ in dwnlds} + settled = [ + typ + for typ, path in paths.items() + if path or (not rerun and log.get(ome, {}).get(typ) == "error") + ] + if len(settled) == len(dwnlds): + preexisting[i] = {t: p for t, p in paths.items() if p} + else: + todwnld_i.append(i) + + if preexisting: + logger.debug(f"{spacer}{len(preexisting)} preexisting download(s)") + for i, paths in preexisting.items(): + for typ, path in paths.items(): + jgi_df.at[i, typ + "_path"] = path + + api_deferred = set() + if todwnld_i: + post_df, api_failed = jgi_dwnld( + jgi_df.loc[todwnld_i].copy(), + output, + user, + pwd, + assembly="fna" in dwnlds, + proteome="faa" in dwnlds, + gff3="gff3" in dwnlds, + masked=masked, + spacer=spacer, + ome_col=ome_col, + deferred=api_deferred, + restore_wait=restore_wait, + defer_tape=True, ) + else: + post_df, api_failed = None, set() - return jgi_df, log, failed - - -def log2df(jgi_df, log, output): + todel, deferred = [], [] + if post_df is not None: + for i, row in post_df.iterrows(): + ome = row[ome_col] + if ome not in log: + log[ome] = {"fna": "na", "gff3": "na", "faa": "na"} + # a file JGI has yet to stage to disk is pending, not failed + unobtained = "pending" if ome in api_deferred else "error" + for typ in dwnlds: + path = row.get(typ + "_path") + if isinstance(path, str) and path: + jgi_df.at[i, typ + "_path"] = path + log[ome][typ] = Path(path).name + logger.debug(f"{spacer}{ome} {typ}: {Path(path).name}") + else: + log[ome][typ] = unobtained + logger.debug(f"{spacer}{ome} {typ}: {unobtained.upper()}") + # JGI metadata may fill in curation fields the MycoCosm table lacks + for col in ("genus", "species", "strain"): + if col in post_df.columns: + jgi_df.at[i, col] = row[col] + log_editor( + log_path, + ome, + ome + + "\t" + + log[ome]["fna"] + + "\t" + + log[ome]["gff3"] + + "\t" + + log[ome]["faa"], + ) + if ome in api_deferred: + deferred.append([ome, jgi_df["version"][i]]) + todel.append(i) + elif ome in api_failed: + failed.append([ome, jgi_df["version"][i]]) + todel.append(i) + + for i in todel: + jgi_df = jgi_df.drop(i) - typ2col = { - "fna": "assemblyPath", - "gff": "gffPath", - "gff3": "gffPath", - "faa": "proteomePath", - } - for ome in list(jgi_df.index): - for typ in log[ome]: - # if log[ome][typ].endswith('.gff.gz'): - # out_file = output + '/gff/' + log[ome][typ] - # col = 'jgi_gff2_path' - # else: - out_file = output + "/" + typ + "/" + log[ome][typ] - jgi_df.at[ome, typ2col[typ]] = out_file + if deferred: + logger.info( + f"{spacer}{len(deferred)} genome(s) awaiting JGI tape restore; " + "they will be retried on the next run" + ) - return jgi_df + return jgi_df, log, failed, deferred def main( @@ -247,6 +290,7 @@ def main( duplicates={}, spacer="\t\t", jgi2ncbi={}, + restore_wait=None, ): if not nonpublished: @@ -255,9 +299,8 @@ def main( ome_col = "assembly_acc" elif "portal" in jgi_df.columns: ome_col = "portal" - name_col = "name" else: - print(spacer + "ERROR: invalid MycoCosm tsv headers", flush=True) + logger.debug(spacer + "invalid MycoCosm tsv headers") sys.exit(3) toDel = [] @@ -281,11 +324,10 @@ def main( jgi_df = jgi_df.drop(failed) jgi_df = jgi_df.reset_index() - print(spacer + "Redundancy check", flush=True) + logger.debug(spacer + "Redundancy check") if isinstance(ref_db, pd.DataFrame): ref_db["index"] = ref_db["assembly_acc"].copy() ref_db = ref_db.set_index("index") - old_len = len(jgi_df) jgi_df, new_ref_db, updates, old_rows = jgi_redundancy_check( ref_db, jgi_df, ome_col=ome_col, jgi2ncbi=jgi2ncbi ) @@ -298,36 +340,13 @@ def main( else: updates.to_csv(f"{output}/jgiUpdates.tsv", sep="\t") # update_check = {i[-1]: i[0:3] for i in updates if i[0]} - print(spacer + "\t" + str(len(jgi_df)) + " genomes to assimilate", flush=True) + logger.debug(spacer + "" + str(len(jgi_df)) + " genomes to assimilate") else: new_ref_db, updates = None, {} - print(spacer + "Logging into JGI", flush=True) - jgi_login(user, pwd) - - if not os.path.exists(output + "/xml"): - os.mkdir(output + "/xml") - - print(spacer + "Retrieving `xml` directories", flush=True) - ome_set, failed, count = set(), [], 0 - for i, row in jgi_df.iterrows(): - error_check, attempt = True, 0 - while error_check != -1 and attempt < 3: - attempt += 1 - error_check = retrieve_xml(row[ome_col], output + "/xml") - if error_check is None: - time.sleep(1) - continue - # elif error_check > 0: - # ome_set.add(row[ome_col]) - elif error_check != -1: - time.sleep(0.3) - if error_check != -1: - eprint(f"{spacer}\t{row[ome_col]} failed to retrieve XML", flush=True) - ome_set.add(row[ome_col]) - + failed = [] log_path = output + "/jgi2db.log" - log = compileLog(log_path) + log = compile_log(log_path) if not rerun: prev_omes = set(jgi_df[ome_col]) for ome in log: @@ -337,9 +356,8 @@ def main( drop_index = list(jgi_df[jgi_df[ome_col] == ome].index)[0] failed.append([ome, jgi_df["version"][drop_index]]) jgi_df = jgi_df.drop(drop_index) - ome_set.add(ome) - print(spacer + "Downloading JGI data", flush=True) + logger.debug(spacer + "Downloading JGI data") dwnlds = [] if assembly: dwnlds.append("fna") @@ -349,45 +367,24 @@ def main( dwnlds.append("faa") for typ in dwnlds: - if not os.path.isdir(output + "/" + typ): - os.mkdir(output + "/" + typ) - if typ == "gff3": - if not os.path.isdir(output + "/gff3"): - os.mkdir(output + "/gff3") - - if all(x in log for x in list(jgi_df[ome_col])) and not rerun: - print(spacer + "\tAll downloaded, rerun off", flush=True) - jgi_df = jgi_df.set_index(ome_col) - jgi_df = log2df(jgi_df, log, output) - jgi_df = jgi_df.reset_index() - else: - for i, row in jgi_df.iterrows(): - ome = row[ome_col] - if ome not in log: - log[ome] = {"fna": "na", "gff3": "na", "faa": "na"} - elif ome not in set(jgi_df[ome_col]): - continue - jgi_df, log, failed = runjgi_dwnld( - jgi_df, - i, - user, - pwd, - ome_set, - ome_col, - output, - log, - log_path, - dwnlds, - failed, - rerun, - repeatmasked, - spacer, - ) - - if os.path.exists("cookies"): - os.remove("cookies") - if os.path.exists(os.path.expanduser("~/.nullJGIdwnld")): - os.remove(os.path.expanduser("~/.nullJGIdwnld")) + if not Path(output + "/" + typ).is_dir(): + Path(output + "/" + typ).mkdir() + + jgi_df, log, failed, deferred = runjgi_dwnld( + jgi_df, + user, + pwd, + ome_col, + output, + log, + log_path, + dwnlds, + failed, + rerun, + repeatmasked, + spacer, + restore_wait=restore_wait, + ) jgi_df = jgi_df.rename( columns={ @@ -399,7 +396,9 @@ def main( ) jgi_df["source"] = "jgi" - for ome_d in failed: # add back the failed entries that were attempted updates + # add back entries whose attempted update did not complete, whether it + # failed outright or is still awaiting a JGI tape restore + for ome_d in failed + deferred: ome = ome_d[0] if ome in old_rows: # if there is an old row to add back new_ref_db = new_ref_db.append(old_rows[ome]) @@ -421,7 +420,7 @@ def main( elif not pd.isnull(row["is public"]) and row["is public"]: jgi_premtdb_df.at[i, "published"] = 1 - return jgi_premtdb_df, new_ref_db.reset_index(), failed + return jgi_premtdb_df, new_ref_db.reset_index(), failed, deferred def cli(): @@ -438,7 +437,7 @@ def cli(): parser.add_argument( "-l", "--login", - help=r'Login file: "\t\n\t"', + help=r'Login file: "\t\n"', ) parser.add_argument("-d", "--database", help="Existing myctools `.db` to reference") parser.add_argument( @@ -486,6 +485,7 @@ def cli(): ) args = parser.parse_args() + setup_logging(verbose=getattr(args, "verbose", False)) args_dict = { "Preexisting db": args.database, @@ -501,28 +501,28 @@ def cli(): output = os.path.abspath(args.output) else: output = start_time.strftime("%Y%m%d") + "_jgi2db" - if not os.path.isdir(output): - os.mkdir(output) + if not Path(output).is_dir(): + Path(output).mkdir() if args.login: with open(args.login, "r") as raw: prep = raw.read() data = [x.split("\t") for x in prep.split("\n")] - user = data[0][0] - pwd = data[0][1] - email = data[1][0] apikey = None - if len(data[1]) > 1: - if data[1][1] != "": - apikey = data[1][1] + if len(data) > 1 and data[1]: + # NCBI API key is the last field of the second line; a legacy + # leading NCBI email column (now unused) is tolerated + api_field = data[1][-1] + if api_field != "": + apikey = api_field ref_db = db2df(format_path(args.database)) jgi_df = main(args.mycocosm, refdb, output) df2db(jgi_df, output + "/new.db") - print( - "\nSuccess! " - + str(len(jgi_df), flush=True) + logger.debug( + "Success! " + + str(len(jgi_df)) + " added to database\n \ Run updateDB to confirm and finish update." ) diff --git a/mycotools/utils/ncbi2db.py b/mycotools/utils/ncbi2db.py index bcaecec..255372a 100755 --- a/mycotools/utils/ncbi2db.py +++ b/mycotools/utils/ncbi2db.py @@ -7,19 +7,14 @@ # NEED TO EDIT REDUNDANCY CHECK TO REFERENCE QUERIED ASSEMBLY ACCESSIONS FROM # BIOSAMPLES +import logging import os import re -import sys import copy -import argparse -import subprocess -import numpy as np -import pandas as pd from datetime import datetime -from mycotools.lib.kontools import intro, outro, eprint -from mycotools.lib.dbtools import db2df, df2db, primaryDB -from mycotools.ncbiDwnld import main as ncbi_dwnld -from mycotools.predb2mtdb import main as predb2mtdb +from mycotools.download.ncbi import main as ncbi_dwnld + +logger = logging.getLogger(__name__) def redundancy_check(db, ncbi_df, ass_acc, duplicates={}): @@ -125,7 +120,7 @@ def main( duplicates={}, check_MD5=True, spacer="\t\t", - fallback=False, + chunk=25, ): os.chdir(out_dir) @@ -157,7 +152,7 @@ def main( ncbi_df = ncbi_df.drop(i) ncbi_df = ncbi_df.reset_index() - print(spacer + "Redundancy check", flush=True) + logger.debug(spacer + "Redundancy check") update_check = {} if ref_db is not None: if len(ref_db) > 0: @@ -174,45 +169,27 @@ def main( update_check = {i[-2]: i for i in updates if i[0]} # dict(update_check) = {assembly_accNEW: [organism, ref organism, # old_assembly_acc]} - print( - spacer + "\t" + str(len(ncbi_df)) + " genomes to assimilate", flush=True - ) + logger.debug(spacer + "" + str(len(ncbi_df)) + " genomes to assimilate") if len(ncbi_df) > 0: - print(spacer + "Initializing NCBI acquisition", flush=True) - if fallback: - from mycotools.ncbi_dwnld_fallback import main as ncbi_dwnld_fallback - - ncbi_df, failed = ncbi_dwnld_fallback( - assembly=assem, - proteome=prot, - gff3=gff, - ncbi_df=ncbi_df, - remove=True, - output_path=out_dir, - column=ass_acc, - ncbi_column="assembly", - check_MD5=check_MD5, - spacer="\t\t\t", - ) - else: - ncbi_df, failed = ncbi_dwnld( - assembly=assem, - proteome=prot, - gff3=gff, - ncbi_df=ncbi_df, - remove=True, - output_path=out_dir, - column=ass_acc, - ncbi_column="assembly", - check_MD5=check_MD5, - verbose=True, - spacer="\t\t\t", - ) + logger.debug(spacer + "Initializing NCBI acquisition") + ncbi_df, failed = ncbi_dwnld( + assembly=assem, + proteome=prot, + gff3=gff, + ncbi_df=ncbi_df, + remove=True, + output_path=out_dir, + column=ass_acc, + ncbi_column="assembly", + check_MD5=check_MD5, + verbose=True, + spacer="\t\t\t", + chunk=chunk, + ) - print( - spacer + "\t" + str(len(ncbi_df)) + " entries with assemblies and gffs", - flush=True, + logger.debug( + spacer + "" + str(len(ncbi_df)) + " entries with assemblies and gffs" ) ncbi_df = ncbi_df.rename( columns={ @@ -230,10 +207,7 @@ def main( ] # remove it from potential updates ref_db = ref_db.set_index("assembly_acc") # update ome codes - # try: ncbi_df = ncbi_df.set_index("assembly_acc") - # except KeyError: # no entries - # return ncbi_df, ref_db.reset_index(), failed, duplicates for assembly_acc, update_d in update_check.items(): old_ome = update_d[-1] ncbi_df.at[assembly_acc, "ome"] = old_ome diff --git a/mycotools/utils/og2mycodb.py b/mycotools/utils/og2mycodb.py index f2058e3..fc93845 100755 --- a/mycotools/utils/og2mycodb.py +++ b/mycotools/utils/og2mycodb.py @@ -1,11 +1,11 @@ #! /usr/bin/env python3 -import os import re import sys import multiprocessing as mp -from mycotools.lib.biotools import gff2list, list2gff, gff3Comps +from mycotools.lib.biotools import gff2list, list2gff, gff3_comps from mycotools.lib.kontools import format_path, sys_start +from pathlib import Path def og2dict(orthogroup_file): @@ -24,7 +24,7 @@ def og2dict(orthogroup_file): return ome_ogs -def sortOGtag(ogtag_dict): +def sort_ogtag(ogtag_dict): sort_list = ["K", "P", "U", "C", "O", "F", "G", "S"] sorted_dict = {} @@ -38,7 +38,7 @@ def sortOGtag(ogtag_dict): return sorted_dict -def readOGtag(ogtagData): +def read_ogtag(ogtagData): ogs = ogtagData.split("|") ogtag_dict = {} @@ -49,19 +49,19 @@ def readOGtag(ogtagData): return ogtag_dict -def writeOGtag(ogtag_dict): +def write_ogtag(ogtag_dict): og_str = "" for i in ogtag_dict: og_str += i + ":" + str(ogtag_dict[i]) + "|" return og_str[:-1] -def editOGtag(ogtag_dict, ogtag, og): +def edit_ogtag(ogtag_dict, ogtag, og): ogtag_dict[ogtag] = og new_oginfo = "OG=" - ogtag_dict = sortOGtag(ogtag_dict) + ogtag_dict = sort_ogtag(ogtag_dict) for i in ogtag_dict: new_oginfo += i + ":" + str(ogtag_dict[i]) + "|" new_oginfo = new_oginfo[:-1] @@ -69,7 +69,7 @@ def editOGtag(ogtag_dict, ogtag, og): return new_oginfo -def mycodbOGs(file_path=format_path("$MYCOGFF3/../ogs.tsv"), omes=set()): +def mycodb_ogs(file_path=format_path("$MYCOGFF3/../ogs.tsv"), omes=set()): ogInfo_dict = {} with open(file_path, "r") as raw: @@ -78,13 +78,13 @@ def mycodbOGs(file_path=format_path("$MYCOGFF3/../ogs.tsv"), omes=set()): og_info = line.rstrip().split("\t") gene, ogtag_info = og_info[0], og_info[1] if gene[: gene.find("_")] in omes: - ogtag_dict = readOGtag(ogtag_info) + ogtag_dict = read_ogtag(ogtag_info) ogInfo_dict[gene] = ogtag_dict else: for line in raw: og_info = line.rstrip().split("\t") gene, ogtag_info = og_info[0], og_info[1] - ogtag_dict = readOGtag(ogtag_info) + ogtag_dict = read_ogtag(ogtag_info) ogInfo_dict[gene] = ogtag_dict return ogInfo_dict @@ -104,10 +104,10 @@ def extract_ogs(ogInfo_dict, ogtag): return og2gene, gene2og -def og2mycoDB(ogInfo_dict, omes=set(), file_path=format_path("$MYCOGFF3/../ogs.tsv")): +def og2mycodb(ogInfo_dict, omes=set(), file_path=format_path("$MYCOGFF3/../ogs.tsv")): out_list = [] - if os.path.isfile(file_path): + if Path(file_path).is_file(): with open(file_path, "r") as raw: for line in raw: data = line.rstrip().split("\t") @@ -116,7 +116,7 @@ def og2mycoDB(ogInfo_dict, omes=set(), file_path=format_path("$MYCOGFF3/../ogs.t out_list.append(line.split("\t")) for gene in ogInfo_dict: - out_list.append([gene, writeOGtag(ogInfo_dict[gene])]) + out_list.append([gene, write_ogtag(ogInfo_dict[gene])]) sorted_list = [ "\t".join([str(x) for x in y]) for y in sorted(out_list, key=lambda x: x[0]) @@ -131,14 +131,14 @@ def og2gff(ogs_dict, gff_path, ogtag): for entry in gff: if entry["type"] == "gene": - gene = re.search(gff3Comps()["Alias"], entry["attributes"])[1] + gene = re.search(gff3_comps()["Alias"], entry["attributes"])[1] if gene in ogs_dict: - ogSearch = re.search(gff3Comps()["OG"], entry["attributes"]) + ogSearch = re.search(gff3_comps()["OG"], entry["attributes"]) if ogSearch: - ogtag_dict = readOGtag(ogSearch[1]) - new_oginfo = editOGtag(ogtag_dict, ogtag, ogs_dict[gene]) + ogtag_dict = read_ogtag(ogSearch[1]) + new_oginfo = edit_ogtag(ogtag_dict, ogtag, ogs_dict[gene]) entry["attributes"] = re.sub( - gff3Comps()["OG"], new_oginfo, entry["attributes"] + gff3_comps()["OG"], new_oginfo, entry["attributes"] ) else: if not entry["attributes"].endswith(";"): @@ -149,11 +149,11 @@ def og2gff(ogs_dict, gff_path, ogtag): out.write(list2gff(gff)) -def dbMain(og_file, ogtag): +def db_main(og_file, ogtag): ome_ogs = og2dict(og_file) omes = set(ome_ogs.keys()) try: - ogInfo_dict = mycodbOGs( + ogInfo_dict = mycodb_ogs( file_path=format_path("$MYCOGFF3/../ogs.tsv"), omes=omes ) except FileNotFoundError: @@ -165,10 +165,10 @@ def dbMain(og_file, ogtag): else: ogInfo_dict[gene] = {ogtag: ome_ogs[ome][gene]} - og2mycoDB(ogInfo_dict, omes=omes, file_path=format_path("$MYCOGFF3/../ogs.tsv")) + og2mycodb(ogInfo_dict, omes=omes, file_path=format_path("$MYCOGFF3/../ogs.tsv")) -def dbMainGff(og_file, ogtag, cpus=1): +def db_main_gff(og_file, ogtag, cpus=1): ome_ogs = og2dict(og_file) @@ -190,11 +190,7 @@ def cli(): + "species [S]" ) args = sys_start(sys.argv[1:], usage, 2, files=[sys.argv[1]]) - if len(args) > 2: - cpus = int(args[2]) - else: - cpus = 1 - dbMain(format_path(args[0]), args[1]) + db_main(format_path(args[0]), args[1]) sys.exit(0) diff --git a/pyproject.toml b/pyproject.toml index f3ed868..f0d07f5 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,56 +4,81 @@ build-backend = "setuptools.build_meta" [project] name = "mycotools" -version = "1.0.0" +version = "2.0.0" authors = [{name="Zachary Konkel", email="konkelzach@protonmail.com"}] description = "Comparative genomics automation and standardization software" readme = "README.md" license = {file = "LICENSE"} -requires-python = '>=3.0,<4' -dependencies = ['biopython', 'pandas', 'requests', 'scipy', 'openpyxl', 'tqdm', - 'cryptography', 'ete3', 'pyqt5'] +requires-python = ">=3.9,<4" +# Versions are lower-bounded to the reference `mycotools` micromamba env +# (Python 3.9.21). numpy, clipkit, cogent3, and dna_features_viewer were +# previously undeclared despite being imported by the package. +# +# The MycotoolsDB SQLite backend (mycotools/lib/mtdb_sql.py) adds no dependency: +# it uses the stdlib `sqlite3` module, which every supported CPython ships. +dependencies = [ + "biopython>=1.85", + "pandas>=2.2.3", + "numpy>=2.0.2", + "scipy>=1.13.1", + "requests>=2.32.3", + "openpyxl>=3.1.5", + "tqdm>=4.67.1", + "cryptography>=44.0.0", + "ete3>=3.1.3", + "pyqt5>=5.15.9", + "clipkit>=2.4.1", + "cogent3>=2025.3.22a2", + "dna_features_viewer>=3.1.5", +] classifiers = [ "Programming Language :: Python :: 3", "License :: OSI Approved :: GNU Lesser General Public License v3 or later (LGPLv3+)", "Operating System :: POSIX :: Linux", ] +# Test-only requirements; not installed for normal use. +[project.optional-dependencies] +test = [ + "pytest>=8", +] + [project.urls] "Homepage" = "https://github.com/xonq/mycotools" "Bug Tracker" = "https://github.com/xonq/mycotools/issues" [project.scripts] +# Primary entrypoints: `mtdb` (database lifecycle) and `mycotools` (analysis). mtdb = "mycotools.mtdb:main" -acc2fa = "mycotools.acc2fa:cli" -acc2gbk = "mycotools.acc2gbk:cli" -acc2gff = "mycotools.acc2gff:cli" -acc2locus = "mycotools.acc2locus:cli" -add2gff = "mycotools.add2gff:cli" -annotationStats = "mycotools.annotationStats:cli" -assemblyStats = "mycotools.assemblyStats:cli" -bioreform = "mycotools.bioreform:cli" -coords2fa = "mycotools.coords2fa:cli" -crap = "mycotools.crap:cli" -db2files = "mycotools.db2files:cli" -db2hgs = "mycotools.db2hgs:cli" -db2microsyntree = "mycotools.db2microsyntree:cli" -db2search = "mycotools.db2search:cli" -extract_mtdb = "mycotools.extract_mtdb:cli" -fa2clus = "mycotools.fa2clus:cli" -fa2hmmer2fa = "mycotools.fa2hmmer2fa:cli" -fa2mass = "mycotools.fa2mass:cli" -fa2tree = "mycotools.fa2tree:cli" -fna2faa = "mycotools.fna2faa:cli" -gff2seq = "mycotools.gff2seq:cli" -gff2svg = "mycotools.gff2svg:cli" -jgiDwnld = "mycotools.jgiDwnld:cli" -manage_mtdb = "mycotools.manage_mtdb:cli" -ncbiAcc2fa = "mycotools.ncbiAcc2fa:cli" -ncbiDwnld = "mycotools.ncbiDwnld:cli" -ome2name = "mycotools.ome2name:cli" -predb2mtdb = "mycotools.predb2mtdb:cli" -s2subs = "mycotools.s2subs:cli" -update_mtdb = "mycotools.update_mtdb:cli" +mycotools = "mycotools.cli:main" + +# DEPRECATED flat aliases -- kept for one transition release. Each prints a +# deprecation notice and forwards to the tool's new nested invocation (shown +# alongside). To be REMOVED in a subsequent release. +add2gff = "mycotools.deprecated:add2gff" # mycotools gff add +annotationStats = "mycotools.deprecated:annotationStats" # mycotools stats annotation +assemblyStats = "mycotools.deprecated:assemblyStats" # mycotools stats assembly +bioreform = "mycotools.deprecated:bioreform" # mycotools seq convert +coords2fa = "mycotools.deprecated:coords2fa" # mycotools seq coords +crap = "mycotools.deprecated:crap" # mycotools phylo crap +db2files = "mycotools.deprecated:db2files" # mtdb files +db2hgs = "mycotools.deprecated:db2hgs" # mycotools cluster db +db2microsyntree = "mycotools.deprecated:db2microsyntree" # mycotools phylo synteny +db2search = "mycotools.deprecated:db2search" # mycotools homology db +fa2clus = "mycotools.deprecated:fa2clus" # mycotools cluster fasta +fa2hmmer2fa = "mycotools.deprecated:fa2hmmer2fa" # mycotools homology fasta +fa2mass = "mycotools.deprecated:fa2mass" # mycotools seq mass +fa2tree = "mycotools.deprecated:fa2tree" # mycotools phylo tree +fna2faa = "mycotools.deprecated:fna2faa" # mycotools seq translate +gff2seq = "mycotools.deprecated:gff2seq" # mycotools seq gff +gff2svg = "mycotools.deprecated:gff2svg" # mycotools gff svg +jgiDwnld = "mycotools.deprecated:jgiDwnld" # mycotools download jgi +ncbiDwnld = "mycotools.deprecated:ncbiDwnld" # mycotools download ncbi +ome2name = "mycotools.deprecated:ome2name" # mycotools rename +acc2fa = "mycotools.deprecated:acc2fa" # mtdb accession fa +acc2gff = "mycotools.deprecated:acc2gff" # mtdb accession gff +acc2gbk = "mycotools.deprecated:acc2gbk" # mtdb accession gbk +acc2locus = "mycotools.deprecated:acc2locus" # mtdb accession locus [tool.setuptools.packages.find] diff --git a/test/README.md b/test/README.md index d89e483..f8a499b 100644 --- a/test/README.md +++ b/test/README.md @@ -23,7 +23,7 @@ guide](https://github.com/xonq/mycotools/blob/master/USAGE.md).
### Reconstruct phylogenies and synteny diagrams of the nitrate assimilation gene cluster using the Cluster Reconstruction and Phylogenetic Analysis Pipeline (CRAP): -```acc2locus -a ustbro1_1795 -p 1 | crap -q - -s blastp -d $(mtdb) -c ``` +```mtdb accession locus -a ustbro1_1795 -p 1 | crap -q - -s blastp -d $(mtdb) -c ```
diff --git a/test/integration/README.md b/test/integration/README.md new file mode 100644 index 0000000..fb5f022 --- /dev/null +++ b/test/integration/README.md @@ -0,0 +1,42 @@ +# Live integration tests + +These tests **hit the network** and **download real genome data**. They are the +opposite of the offline suites (`test/unit`, `test/update_mtdb`) and are marked +`@pytest.mark.integration`. + +## Requirements + +- The `mycotools` conda env (pandas, biopython, and the `datasets` executable on + PATH). +- Network access to NCBI (and JGI, when its API works). +- **Stored no-password credentials** at `~/.mycotools/mtdb_credentials.json` + (set via `mtdb-manage -s`). A missing store is a hard **error**, not a skip. + +## Run + +```bash +micromamba run -n mycotools pytest test/integration/ -v -s +``` + +`test_init_primary_db_downloads_jgi_and_ncbi` initializes a throwaway primary +MTDB from a 1-JGI + 1-NCBI subset of `test/ust.mtdb` and asserts genome data +from **both** sources is downloaded and curated (the finished primary `.mtdb` +must contain a `jgi`- and an `ncbi`-sourced row, each with `.fna`/`.faa`/`.gff3` +files). It runs in an **isolated `HOME`** (a temp dir holding only a copy of the +credential file) so the real `~/.mycotools/config.json` and your linked/active +database are never touched. + +## JGI download API + tape restores + +JGI genome downloads use the current **JGI Data Portal API** (see +`mycotools/jgiDwnld.py`): `mycocosm_file_list` (search) → `request_archived_files` +(restore) → `download_files` (zip stream), authorized with a signon session +token. The old `get-directory` XML endpoint was retired by JGI. + +JGI keeps most files in tape archive (`file_status: PURGED`); the first download +of a genome requests a restore to disk, which "typically takes less than an hour +but can take up to a night". To stay deterministic, the test's fixture pre-warms +the restore for its JGI genome (bounded by `JGI_RESTORE_TIMEOUT`). If the restore +has not completed within that bound the test **skips** with a message to rerun +once JGI has staged the data, rather than blocking on tape I/O. Once a genome's +files are on disk they remain downloadable (no wait) until they re-purge. diff --git a/test/integration/conftest.py b/test/integration/conftest.py new file mode 100644 index 0000000..bc95ac4 --- /dev/null +++ b/test/integration/conftest.py @@ -0,0 +1,5 @@ +def pytest_configure(config): + config.addinivalue_line( + "markers", + "integration: live network test (downloads JGI/NCBI data; needs credentials)", + ) diff --git a/test/integration/test_init_download.py b/test/integration/test_init_download.py new file mode 100644 index 0000000..4698e57 --- /dev/null +++ b/test/integration/test_init_download.py @@ -0,0 +1,268 @@ +#! /usr/bin/env python3 +"""LIVE integration test: initialize a primary MycotoolsDB from test/ust.mtdb. + +Unlike the offline suites (test/unit, test/update_mtdb), this test actually +downloads JGI (MycoCosm) and NCBI genome data and gathers NCBI taxonomy, using +the *no-password* credential store (``~/.mycotools/mtdb_credentials.json``, see +dbtools.store_login). It therefore requires: + + * network access to JGI and NCBI, + * the `datasets` executable on PATH (ships in the `mycotools` conda env), and + * stored, accessible NCBI/JGI credentials. + +Per request, a missing credential store is a hard ERROR (not a skip): the test +raises so the absence is loud. + +JGI downloads use the current JGI Data Portal API (see mycotools.jgiDwnld). JGI +archives most files to tape (file_status PURGED); a first download restores them +to disk, which "typically takes less than an hour but can take up to a night". +To keep this test deterministic and fast, its fixture pre-warms the restore for +the JGI genome (bounded); if the restore has not completed within that bound the +test skips (rerun once JGI has staged the data) rather than block on tape I/O. + +Isolation: the run happens in a throwaway ``HOME`` that contains only a copy of +the credential file - never the real ``~/.mycotools/config.json`` - so the +user's linked/active database is never modified. update_mtdb writes its new +config into the temp HOME, which is deleted on teardown. + +Run: + + micromamba run -n mycotools pytest test/integration/ -v -s + +To bound runtime, only a small subset of ust.mtdb is initialized (see +N_PER_SOURCE); one JGI + one NCBI entry still exercises the full download and +curation paths for both sources. +""" +import os +import shutil +import subprocess +import sys +from pathlib import Path + +import pytest + +from mycotools.lib.dbtools import read_plain_login + +REPO_ROOT = Path(__file__).resolve().parents[2] +UST_MTDB = REPO_ROOT / "test" / "ust.mtdb" +CREDS_PATH = Path.home() / ".mycotools" / "mtdb_credentials.json" + +# entries per source (jgi/ncbi) to pull from ust.mtdb into the reference; small +# by default so the live download stays quick. +N_PER_SOURCE = 1 + +# column indices in an .mtdb row (0-based): ome, genus, species, strain, +# taxonomy, version, source, biosample, assembly_acc, ... +_SOURCE_COL = 6 +_ASSEMBLY_ACC_COL = 8 + +# how long to wait for JGI to restore archived (on-tape) files before skipping +JGI_RESTORE_TIMEOUT = 1200 + + +def require_credentials(): + """Return (ncbi_email, ncbi_api, jgi_email, jgi_pwd) from the no-password + store, raising if the store or any required field is absent.""" + if not CREDS_PATH.is_file(): + raise RuntimeError( + f"no-password credential store not found at {CREDS_PATH}. Store " + "credentials first with `mtdb manage --store` (or `mtdb m -s`)." + ) + ncbi_email, ncbi_api, jgi_email, jgi_pwd = read_plain_login(str(CREDS_PATH)) + required = {"ncbi_email": ncbi_email, "jgi_email": jgi_email, "jgi_pwd": jgi_pwd} + missing = sorted(name for name, val in required.items() if not val) + if missing: + raise RuntimeError( + f"credential store {CREDS_PATH} is missing required field(s): {missing}" + ) + return ncbi_email, ncbi_api, jgi_email, jgi_pwd + + +def _iter_rows(source=None): + """Yield the tab-split fields of each data row in ust.mtdb, optionally + filtered to a single source (jgi/ncbi).""" + for line in UST_MTDB.read_text().splitlines(): + if not line.strip() or line.startswith("#"): + continue + fields = line.split("\t") + if source is None or fields[_SOURCE_COL].strip().lower() == source: + yield fields + + +def build_reference_subset(dest: Path, n_per_source: int = N_PER_SOURCE) -> Path: + """Write a small reference .mtdb containing n JGI + n NCBI entries taken from + ust.mtdb (preserving its exact columns).""" + rows = [] + for src in ("jgi", "ncbi"): + rows.extend(["\t".join(f) for f in list(_iter_rows(src))[:n_per_source]]) + assert any("\tjgi\t" in r or r.split("\t")[_SOURCE_COL] == "jgi" for r in rows) + assert any(r.split("\t")[_SOURCE_COL] == "ncbi" for r in rows) + dest.write_text("\n".join(rows) + "\n") + return dest + + +def jgi_portal_ids(n_per_source: int = N_PER_SOURCE): + """The JGI portal ids (assembly_acc) that build_reference_subset will use.""" + return [ + f[_ASSEMBLY_ACC_COL].strip() for f in list(_iter_rows("jgi"))[:n_per_source] + ] + + +def warm_jgi_restore(portal_ids, jgi_email, jgi_pwd, timeout=JGI_RESTORE_TIMEOUT): + """Ensure the assembly + gff3 for each JGI portal id are RESTORED (on disk), + requesting a tape restore and polling if needed. Dogfoods the jgiDwnld API + helpers so the warmed selection matches what update_mtdb will download. + Returns the set of portal ids whose essential files are ready.""" + from mycotools.jgiDwnld import ( + jgi_api_login, + search_organism, + select_file, + request_restore, + poll_restore, + _mycocosm_ids, + _is_restored, + ) + + session, token = jgi_api_login(jgi_email, jgi_pwd) + ready = set() + for portal_id in portal_ids: + org, files = search_organism(session, portal_id) + if not org: + continue + selected = [ + f + for f in ( + select_file(files, "fna", masked=True), + select_file(files, "gff3", masked=True), + ) + if f is not None + ] + if len(selected) < 2: # need both an assembly and an annotation + continue + purged = [f["_id"] for f in selected if not _is_restored(f)] + if purged: + status_url = request_restore( + session, + token, + _mycocosm_ids( + org.get("id"), + (org.get("top_hit") or {}).get("_id"), + org.get("mycocosm_portal_id") or portal_id, + purged, + ), + ) + if not poll_restore(session, status_url, timeout=timeout, interval=20): + continue + ready.add(portal_id) + return ready + + +@pytest.fixture +def isolated_home(tmp_path): + """A throwaway HOME containing only a copy of the real credential store.""" + require_credentials() # hard error if creds are missing + home = tmp_path / "home" + (home / ".mycotools").mkdir(parents=True) + dst = home / ".mycotools" / "mtdb_credentials.json" + shutil.copy(CREDS_PATH, dst) + os.chmod(dst, 0o600) + return home + + +def run_update_mtdb(home: Path, init_dir: Path, reference: Path, timeout=1800): + """Invoke `update_mtdb --init --reference ` in an + isolated HOME, with MYCODB stripped so no linked DB is inherited.""" + env = dict(os.environ) + env["HOME"] = str(home) + for var in ("MYCODB", "MYCOFNA", "MYCOFAA", "MYCOGFF3"): + env.pop(var, None) + return subprocess.run( + [ + sys.executable, + "-m", + "mycotools.mtdb.update", + "--init", + str(init_dir), + "--reference", + str(reference), + ], + stdin=subprocess.DEVNULL, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + timeout=timeout, + env=env, + ) + + +def read_primary_sources(init_dir: Path): + """Return the list of `source` values in the produced primary .mtdb.""" + mtdb_files = list(init_dir.glob("**/mtdb/*.mtdb")) + assert mtdb_files, "no .mtdb produced" + primary = max(mtdb_files, key=lambda p: p.stat().st_size) + sources = [] + for line in primary.read_text().splitlines(): + if not line.strip() or line.startswith("#"): + continue + fields = line.split("\t") + if len(fields) > _SOURCE_COL: + sources.append(fields[_SOURCE_COL].strip().lower()) + return primary, sources + + +@pytest.mark.integration +def test_credentials_present(): + """The no-password credential store exists and has NCBI + JGI logins.""" + ncbi_email, ncbi_api, jgi_email, jgi_pwd = require_credentials() + assert "@" in ncbi_email + assert "@" in jgi_email + assert jgi_pwd + + +@pytest.mark.integration +def test_init_primary_db_downloads_jgi_and_ncbi(isolated_home, tmp_path): + """Initialize a primary MTDB from a ust.mtdb subset (1 JGI + 1 NCBI) and + confirm genome data from *both* sources is downloaded and curated. + + JGI downloads use the current JGI Data Portal API; the fixture warms the + tape restore first so the download is deterministic. Both the JGI and NCBI + genomes must appear in the finished primary database, each with sequence and + annotation files, and the init must never die on an XML-parse error.""" + _, _, jgi_email, jgi_pwd = require_credentials() + + portal_ids = jgi_portal_ids() + ready = warm_jgi_restore(portal_ids, jgi_email, jgi_pwd) + if set(portal_ids) - ready: + pytest.skip( + "JGI has not finished restoring " + f"{sorted(set(portal_ids) - ready)} from tape; rerun once staged." + ) + + reference = build_reference_subset(tmp_path / "reference_subset.mtdb") + init_dir = isolated_home / "mtdb_init" # non-existent -> becomes the DB root + + result = run_update_mtdb(isolated_home, init_dir, reference) + + # surface the full log on failure so download / curation errors are visible + assert result.returncode == 0, ( + f"update_mtdb --init exited {result.returncode}\n" + f"----- output -----\n{result.stdout}" + ) + + # regression guard: the JGI directory data must never crash the run + assert "Traceback (most recent call last)" not in result.stdout, result.stdout + assert "ParseError" not in result.stdout, result.stdout + + # the finished primary database contains BOTH sources + primary, sources = read_primary_sources(init_dir) + assert sources, f"primary MTDB is empty\n{result.stdout}" + assert "jgi" in sources, f"JGI genome missing from primary DB\n{result.stdout}" + assert "ncbi" in sources, f"NCBI genome missing from primary DB\n{result.stdout}" + + # downloaded + curated sequence/annotation data (>=2 genomes: 1 jgi + 1 ncbi) + fna = list(init_dir.glob("**/data/fna/*.fna")) + faa = list(init_dir.glob("**/data/faa/*.faa")) + gff3 = list(init_dir.glob("**/data/gff3/*.gff3")) + assert len(fna) >= 2, f"expected >=2 genome assemblies (.fna)\n{result.stdout}" + assert len(faa) >= 2, f"expected >=2 proteomes (.faa)\n{result.stdout}" + assert len(gff3) >= 2, f"expected >=2 annotations (.gff3)\n{result.stdout}" diff --git a/test/mycotools_workshop.md b/test/mycotools_workshop.md index 1342e2e..99172a2 100644 --- a/test/mycotools_workshop.md +++ b/test/mycotools_workshop.md @@ -84,15 +84,15 @@ jgiDwnld -i Ustbr1 -a -g This script will output a `predb` file that is ready for assimilating into the database. If you wanted to add your own genomes, you would fill out one of -these files manually by generating a blank copy via `mtdb predb2mtdb > predb.tsv`, +these files manually by generating a blank copy via `mtdb predb > predb.tsv`, then running the following commands as we will here: ```bash -# curate the data via predb2mtdb +# curate the data via predb mtdb p Ustbr1.predb.tsv # add the curated data to the primary MTDB -mtdb u -a predb2mtdb_/predb2mtdb.mtdb +mtdb u -a predb_/predb.mtdb ``` Now we can check if the file was added by querying the genome code from the @@ -196,8 +196,8 @@ have time, so let's work with a subset by copying them to a new folder: mkdir sco_202405 # copy the top three SCOs -for i in $(ls db2hgs_/single_copy_genes/ | head -3) - do cp db2hgs_/single_copy_genes/$i sco_202405/ +for i in $(ls cluster_db_/single_copy_genes/ | head -3) + do cp cluster_db_/single_copy_genes/$i sco_202405/ done # run the tree building pipeline @@ -205,7 +205,7 @@ fa2tree -i sco_202405/ --partition ``` When complete, we will open the `concatenated.nex.contree` file in the -resulting `fa2tree_` directory in FigTree, which is the consensus +resulting `phylo_tree_` directory in FigTree, which is the consensus tree with 1000 ultrafast bootstrap replicates. What you will note is that the tips are labeled with the ome code - but we @@ -213,8 +213,8 @@ probably want to see the actual genus, species, and strain names, right?! Let's convert the phylogenomic tree from genome code tips to full names: ```bash -ome2name fa2tree_/concatenated.nex.contree o \ - > fa2tree_/full_name.newick +ome2name phylo_tree_/concatenated.nex.contree o \ + > phylo_tree_/full_name.newick ``` Go ahead and open this one in FigTree, and let's glance at how well supported @@ -231,11 +231,11 @@ assimilation in our dataset. First, we need to identify homologs of this gene across our database. We will do this by implementing a BLAST search of the protein sequence. We obtain the -protein sequence using a handy command, `acc2fa`. +protein sequence using a handy command, `mtdb accession fa`. ```bash # extract the protein accession of interest -acc2fa -a ustbro1_1795 > ustbro1_1795.faa +mtdb accession fa -a ustbro1_1795 > ustbro1_1795.faa # run a blast search on this gene against the primary MTDB db2search -a blastp -q ustbro1_1795.faa -e 2 @@ -249,10 +249,10 @@ building at the cost of some quality: ```bash # move the phylogenomic directory -mv fa2tree_ phylogenomic_/ +mv phylo_tree_ phylogenomic_/ # run the single gene phylo -fa2tree -i db2search_/fastas/ustbro1_1795.search.fa -f +fa2tree -i homology_db_/fastas/ustbro1_1795.search.fa -f ``` Now, we can view this tree in FigTree. @@ -276,7 +276,7 @@ locus, then inputting it into the CRAP pipeline: ```bash # extract a locus of interest, and store in a file -acc2locus -a ustbro1_1795 -p 1 > nitrate_cluster.txt +mtdb accession locus -a ustbro1_1795 -p 1 > nitrate_cluster.txt # run the CRAP analysis crap -q nitrate_cluster.txt -s blastp @@ -311,7 +311,7 @@ mkdir clinker_ Then generate GenBanks of each locus file using some basic BASH scripting: ```bash -for i in crap_/loci/*txt; do o=$(basename ${i} .txt); acc2gbk -i ${i} +for i in crap_/loci/*txt; do o=$(basename ${i} .txt); mtdb accession gbk -i ${i} > clinker_/${o}.gbk; done ``` diff --git a/test/unit/README.md b/test/unit/README.md new file mode 100644 index 0000000..d8e14b4 --- /dev/null +++ b/test/unit/README.md @@ -0,0 +1,51 @@ +# mycotools unit-test scaffold + +Offline, quick unit tests for the mycotools scripts. No JGI/NCBI login, no +downloads, no genome assimilation — every test exercises argument handling, +argparse surfaces, or login-free helper functions. + +## Run + +```bash +micromamba run -n mycotools pytest test/unit/ -q # this scaffold +micromamba run -n mycotools pytest test/ -q # scaffold + update_mtdb suite +``` + +(~9 s.) + +## Layout + +| File | What it is | +| --- | --- | +| `conftest.py` | Shared fixtures: `offline_env` (autouse; strips `MYCODB`, sets dummy `MYCO*` prefixes), path fixtures (`ust_mtdb`, `reference_mtdb`, `repo_root`, `test_data_dir`), and a `run_cli` subprocess helper. | +| `test_cli_smoke.py` | Data-driven **import** + **`--help` doesn't crash** checks for every entry-point script. The backbone: one parametrized case per script. | +| `test_lib_biotools.py` | A concrete, worked example of deeper unit tests (pure `fa2dict`/`dict2fa` parsers) — the pattern to copy for real logic. | +| `_template.py` | Copy to `test_