<commit_before>"""
Read a blast file that has gi| numbers and add the taxonomy to the end of the reads. We will add

kingdom phylum class order family genus species
"""

import os
import sys

import argparse
import re
from taxon import taxonomy_hierarchy_as_list, connect_to_db, read_acc_tax_id

default_tax_dir = os.path.join(os.environ['HOME'], "ncbi", "taxonomy", "current")
parser = argparse.ArgumentParser(description="Add taxonomy to a blast output file")
parser.add_argument('-b', help='blast ouput file', required=True)
parser.add_argument('-o', help="Output file (blast results + taxonomy", required=True)
parser.add_argument('-p', help='blast program (blastn, blastp, or blastx (default)', default='blastx')
parser.add_argument('-q', help='Look for taxid of query (default is to use database', action='store_true')
parser.add_argument('-t', help=f'taxonomy directory (default={default_tax_dir})', default=default_tax_dir)

args = parser.parse_args()

bl_col = 1
if args.q:
    bl_col = 0

dbtype='nucl'
if args.p == 'blastx' or args.p == 'blastp':
    dbtype='prot'

if not os.path.exists(os.path.join(args.t, "taxonomy.sqlite3")):
    print(f"FATAL: taxonomy.sqlite3 not found in {args.t}. Please load the database first", file=sys.stderr)
    sys.exit(1)

if not os.path.exists(os.path.join(args.t, "accession2taxid")):
    print(f"FATAL: accession2taxid  not found in {args.t}. Please download that data first", file=sys.stderr)
    sys.exit(1)

print("Reading taxonomy", file=sys.stderr)
acc2tax = read_acc_tax_id(dbtype=dbtype, tax_dir=args.t, verbose=args.v)
print("Read taxonomy", file=sys.stderr)

conn = connect_to_db(os.path.join(args.t, "taxonomy.sqlite3"))

no_result = ['', '', '', '', '', '', '', '']

results = {}
with open(args.b, 'r') as f, open(args.o, 'w') as out:
    for l in f:
        p=l.strip().split("\t")
        m = []
        if p[bl_col] in acc2tax:
            tid = acc2tax[p[bl_col]]
            tax = taxonomy_hierarchy_as_list(conn=conn, tid=tid, verbose=args.v)
            print("\t".join(map(str, p+[str(tid)]+tax)), file=out)
        else:
            if args.v:
                print(f"Warning: No taxid for {p[bl_col]} found", file=sys.stderr)
            print("\t".join(map(str, p+no_result)), file=out)

<commit_msg>Streaming blast results, one line at a time<commit_after>"""
Read a blast file that has gi| numbers and add the taxonomy to the end of the reads. We will add

kingdom phylum class order family genus species
"""

import os
import sys

import argparse
import re
from taxon import taxonomy_hierarchy_as_list, connect_to_db, read_acc_tax_id
from roblib import stream_blast_results

default_tax_dir = os.path.join(os.environ['HOME'], "ncbi", "taxonomy", "current")
parser = argparse.ArgumentParser(description="Add taxonomy to a blast output file")
parser.add_argument('-b', help='blast ouput file', required=True)
parser.add_argument('-o', help="Output file (blast results + taxonomy", required=True)
parser.add_argument('-p', help='blast program (blastn, blastp, or blastx (default)', default='blastx')
parser.add_argument('-q', help='Look for taxid of query (default is to use database', action='store_true')
parser.add_argument('-t', help=f'taxonomy directory (default={default_tax_dir})', default=default_tax_dir)

args = parser.parse_args()

bl_col = 1
if args.q:
    bl_col = 0

dbtype='nucl'
if args.p == 'blastx' or args.p == 'blastp':
    dbtype='prot'

if not os.path.exists(os.path.join(args.t, "taxonomy.sqlite3")):
    print(f"FATAL: taxonomy.sqlite3 not found in {args.t}. Please load the database first", file=sys.stderr)
    sys.exit(1)

if not os.path.exists(os.path.join(args.t, "accession2taxid")):
    print(f"FATAL: accession2taxid  not found in {args.t}. Please download that data first", file=sys.stderr)
    sys.exit(1)

print("Reading taxonomy", file=sys.stderr)
acc2tax = read_acc_tax_id(dbtype=dbtype, tax_dir=args.t, verbose=args.v)
print("Read taxonomy", file=sys.stderr)

conn = connect_to_db(os.path.join(args.t, "taxonomy.sqlite3"))

no_result = ['', '', '', '', '', '', '', '']

results = {}

with open(args.o, 'w') as out:
    for br in stream_blast_results(args.b):
        tid = None
        prot = br.db
        if args.q:
            prot = br.query
        if prot in acc2tax:
            tid = acc2tax[prot]
        if not tid:
            if args.v:
                print(f"Warning: No taxid for {prot} found", file=sys.stderr)
            print("\t".join(map(str, br+no_result)), file=out)

        tax = taxonomy_hierarchy_as_list(conn=conn, tid=tid, verbose=args.v)
        print("\t".join(map(str, br+[str(tid)]+tax)), file=out)

{'fraction_non_alphanumeric': 0.15064655172413793, 'fraction_numerical': 0.005172413793103448, 'mean_word_length': 4.1452328159645235, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6412310', 'n_tokens_mistral': 1646, 'n_tokens_neox': 1612, 'n_words': 415}