<gh_stars>1-10
#!/usr/bin/env python

# proqueryote
# UNH MBCBS 913 Spring 2020
# <NAME>


import os
import sys
import datetime
import argparse
import re
import pandas as pd
from typing import List
from pathlib import Path
from collections import namedtuple
Values = List[str]

'''
=========================================================
GLOBAL CONFIGURATION, INITIALIZATION, AND SUPPORT METHODS
=========================================================
'''

# CACHE FILE LOCATIONS
# Cache is stored in hidden folder in user's home folder.
# Cache contains database files used to process user queries.
# (Actual sequence files are not cached.)
HOME = str(Path.home())
CACHE = HOME + '/.proqueryote/'
PROKARYOTES = CACHE + 'prokaryotes.txt'
NODES_FILE = CACHE + 'nodes.dmp'
NAMES_FILE = CACHE + 'names.dmp'
AUGMENTED = CACHE + 'prokaryotes_taxonomic.txt'
TAXDUMP_FILENAME = 'taxdump.tar.gz'
TAXDUMP = CACHE + '/' + TAXDUMP_FILENAME
CACHE_FILES = [PROKARYOTES, NAMES_FILE, NODES_FILE, AUGMENTED]
CACHE_URLS = [
  'ftp://ftp.ncbi.nlm.nih.gov/genomes/GENOME_REPORTS/prokaryotes.txt',
  'https://ftp.ncbi.nlm.nih.gov/pub/taxonomy/' + TAXDUMP_FILENAME
]

# SPECIFY SEQUENCE FILE TYPES TO DOWNLOAD
FAA_EXTENSION = '/*protein.faa.gz'
FNA_EXTENSION = '/*_cds_from_genomic.fna.gz'

# ORDERED LISTING OF ALL RANKS USED BY NODES.DMP (FROM TAXDUMP.TAR.GZ FROM NCBI)
# This ordered list is necessary for building the taxonomy for each initial taxID listed in prokaryotes.txt
RANK_ORDER = ['superkingdom', 'kingdom', 'superphylum', 'phylum', 'subphylum', 'superclass', 'class', 'subclass', 'infraclass', 'order', 'superfamily', 'family', 'subfamily', 'genus', 'subgenus', 'species group', 'species subgroup', 'species', 'subspecies']

# RANKS FOR WHICH TO ADD COLUMNS TO PROKARYOTES TABLE
# These are the specific taxonomic ranks whose informamtion we wanted
# added as new columns in the prokaryotes.txt table.
def RANKS():
   return {'species':'','genus':'','family':'','phylum':''}

# FOR COMMAND LINE HELP INFO
epilog = """
See query file format specification in associated README.
"""

# SET UP COMMAND LINE ARGUMENTS
parser = argparse.ArgumentParser(prog="column_search", usage="column_search <queries_file>", description="Searches prokaryotes_taxanomic.txt file in working directory for all rows matching queries in given <queries_file>.", epilog=epilog, formatter_class=argparse.RawTextHelpFormatter)

parser.add_argument("queries_file", type=str)
parser.add_argument('--verbose', '-v', action='count', default=0,)
parser.add_argument('--fna', action='store_const', const=True, default=False)
args = parser.parse_args()

# PRINT A MESSAGE TO STANDARD OUTPUT PER VERBOSITY LEVEL
# Only prints messages with level equal to or less than
# verbosity level specified by user on command line.
def log(msg, level=0):
   if (args.verbose >= level): print(msg)

'''
==============
CACHE HANDLING
==============
'''

# TEST CACHE PRESENCE
# Returns False if any of the cache files, or the cache directory, is missing).
def testCachePresence():
  log('testing cache presence', 2)
  if (os.path.isdir(CACHE)):
    log('testCachePresence: cache directory present', 3)
    for each in CACHE_FILES:
      if not os.path.exists(each):
        log(f'testCachePresence: {each} is missing from cache', 3)
        return False
      else:
        log(f'testCachePresence: {each} file is present', 3)
    log('testCachePresence: cache found to be present. OK.', 3)
    return True
  else:
    log('testCachePresence: cache directory missing', 3)
    return False

# UPDATE CACHE
# Removes current cache, if present, and downloads entirely new cache.
def updateCache():
  log('Updating local cache.')
  # remove old cache, create new
  # (will need to make this more fine-grained if we ever
  # store config files in the hidden folder)
  os.system(f'rm -r {CACHE} 2> /dev/null')
  os.mkdir(CACHE)
  
  # download prokaryotes and nodes DB and extract
  working_directory = os.getcwd()
  os.chdir(CACHE)
  log('Downloading resources for local cache', 1)
  for resource in CACHE_URLS:
    log(f'{resource}', 2)
    os.system(f'wget --quiet {resource}')
  os.system(f'tar -z -xf {TAXDUMP}')
  os.chdir(working_directory)

  # created augmented version of table, with new taxonomic rank columns
  taxonomifier = Taxonomifier()
  taxonomifier.produceTaxonomic()


'''
================================================================
BUILD TAXONOMIC INFORMATION AND ADD TO NEW PROKARYOTES.TXT TABLE
================================================================
'''

# PARENT AND RANK
# A convenience type: for a given tax ID, store its rank, and the taxID of its parent.
ParentAndRank = namedtuple('ParentAndRank', 'parent rank')

# TAXONOMIFIER
# Class to organize methods and data for adding new taxonomic rank columns to prokaryotes.txt
class Taxonomifier(object):
  def __init__(self):
    log(f'Preparing to augment table with additional taxonomic rank columns', 1)
    log(f'Loading local cache', 2)
    log(f'Loading {PROKARYOTES}', 3)
    self.proks = open(PROKARYOTES, 'r').readlines()
    self.nodes = self.getNodesMap()
    self.names = self.getNamesMap()
    log(f'Read in {len(self.proks)} lines from {PROKARYOTES}', 3)

  # GET NAMES MAP
  # Load the names.dmp database to memory, for fast querying,
  # as a dictionary of taxID:name key:value pairs.
  @staticmethod
  def getNamesMap():
    log(f'Loading {NAMES_FILE}', 3)
    names = dict()
    with open(NAMES_FILE, 'r') as names_f:
        names_d = names_f.readlines()
        for record in names_d:
          cells = re.split(r'\t\|\t', record)
          taxid = cells[0]
          name = cells[1]
          name_type = cells[3]
          if (name_type.startswith('scientific name')):
              names.setdefault(int(taxid), name)
    return names

  # GET NODES MAP
  # Load the names.dmp database to memory, for fast querying,
  # as a dictionary of taxID:ParentAndRank key:value pairs.
  # (ParentAndRank is the tuple type defind above.)
  @staticmethod
  def getNodesMap():
    log(f'Loading {NODES_FILE}', 3)
    nodes = dict()
    with open(NODES_FILE, 'r') as nodes_f:
        nodes_d = nodes_f.readlines()
        for record in nodes_d:
          cells = re.split(r'\t\|\t', record)
          this_taxid = cells[0]
          this_parent_taxid = cells[1]
          this_rank = cells[2]
          parent_rank = ParentAndRank(int(this_parent_taxid), this_rank)
          nodes.setdefault(int(this_taxid),parent_rank)
    return nodes

  def getName(self, taxid):
    """
    Retrieves scientific name corresponding to given taxid.
    """
    return self.names[int(taxid)]

  def getParent(self, taxid):
    """
    Retrieves taxid of parent of given taxid.
    """
    return self.nodes[int(taxid)].parent

  def getRank(self, taxid):
    """
    Retrieves rank of given taxid.
    """
    return self.nodes[int(taxid)].rank

  @staticmethod
  def getTaxID(proks_record):
    """
    Parse TaxID field from raw string of given row from prokaryotes.txt
    """
    return proks_record.split('\t')[1]

  def getOneTaxonomy(self, taxid):
    """
    Searches nodes.dmp for desired ranks, pulling corresponding names from names.dmp.
    This method builds the string that can be appended to a given row from prokaryotes.txt.

    :return: A TSV string: "<species>\t<genus>\t<family>\t<phylum>"
    """
    ranks = RANKS()
    searchTaxID = str(taxid)
    log(f'getOneTaxonomy: taxid {taxid}', 4)
    # continue searching until we've found an entry for each desired taxonomic rank
    while ('' in ranks.values()):
        log(f'getOneTaxonomy: searching for {searchTaxID}', 4)
        if (searchTaxID == '1'):
          log(f'getOneTaxonomy: arrived at root, curtailing search with extant results: {ranks}', 4)
          break
        try:
          this_parent_taxid = self.getParent(searchTaxID)
        except KeyError:
          log(f'getOneTaxonomy: caught KeyError raised by getParent: TaxID {searchTaxID} is not listed in {args.nodes_file}.', 4)
          break
        this_rank = self.getRank(searchTaxID)
        log(f'getOneTaxonomy: {searchTaxID} is a {this_rank} with parent {this_parent_taxid}', 4)
        # TODO: may want to alter this to confirm that ranks[this_rank] == '',
        # to prevent over-writing of multiple entries exist in nodes.dmp, not all correct?:
        if (this_rank in ranks):
          name = self.getName(searchTaxID)
          ranks[this_rank] = name
          log(f'getOneTaxonomy: added {searchTaxID} : {name} as {this_rank}. we have thus far found: {ranks}', 4)
        searchTaxID = str(this_parent_taxid)
    log(f'getOneTaxonomy: found {ranks}', 4)
    return(ranks)

  def produceTaxonomic(self):
    """
    Create a new prokaryotes_taxonomic.txt table file, with additional taxonomic rank columns.
    """
    log('Building new table', 1)
    # costruct new header row
    old_header = self.proks[0].rstrip()
    new_cols = RANKS().keys()
    log('old header:\n' + old_header, 3)
    log('new columns:\n' + str(new_cols), 4)
    new_header = old_header
    for rank in RANK_ORDER:
      if (rank in new_cols):
        new_header += '\t' + rank.title()
    new_header += '\n'
    log('new header:\n' + new_header, 3)

    # start writing to new file
    with open(AUGMENTED, 'w') as out_f:
      # write new header row
      log(f'About to write new header to new, augmented table:\n{new_header}', 4)
      out_f.write(new_header)
      
      # construct new entries
      for record in self.proks[1:]:
        log('old row:\n' + record, 4)
        taxid = self.getTaxID(record)
        taxonomy = self.getOneTaxonomy(taxid)
        new_record = record.rstrip()
        for rank in RANK_ORDER:
          if (rank in taxonomy.keys()):
              new_record += '\t' + taxonomy[rank]
        new_record += '\n'
        log('new record:\n' + new_record, 4)
        out_f.write(new_record)
    log(f'New table with additional taxonomic columns:\n\t{AUGMENTED}', 2)





'''
=======================================================================================================
PARSE AND PROCESS QUERIES, NOTIFY USER OF NUMBER OF RESULTS, AND PROMPT TO DOWNLOAD SELECTED SEQUENCES.
=======================================================================================================
'''

class Criterion(object):
  '''
  Specifies acceptable values for a given column, for a given query.
  '''
  def __init__(self, column: str, values: Values):
    self.column = column
    self.values = values
  def __repr__(self):
    return f'{self.column}: {self.values}'

class Query(object):
  '''
  Specifies all criteria (in form of Criterion objects) for a given query.
  '''
  def __init__(self):
    self.criteria = list()
  
  def addCriterion(self, criterion: Criterion):
    self.criteria.append(criterion)
  
  def __repr__(self):
    return 'Query:\n' + str(self.criteria)

def parseQueries(queries_file):
  '''
  Parse a query set file, in preparation for processing/executing the queries contained therein.
  '''
  log(f'Parsing query file: {queries_file}', 1)
  queries = list()
  with open(queries_file, 'r') as queries_f:
    for line in queries_f.readlines():
      log("  starting loop", 3)
      line = line.rstrip()
      # skip blank and comment lines
      if (line == '' or line[0] == '#'):
        log("  found empty line or comment; skipping", 3)
        continue
      # start a new query
      elif (line == 'query'):
        log("  found new query", 3)
        queries.append(Query())
      # append criterion to current query
      else:
        log("  adding to criterion to current query", 3)
        delimited = re.split(' ', line)
        queries[len(queries) - 1].addCriterion(Criterion(delimited[0],delimited[1:]))        
  log(f'parsed queries:\n{queries}', 3)
  return queries


class Data(object):
  '''
  Class organizing the data table to be queried, the parsed
  queries, and methods for processing queries against the data table.
  '''
  def __init__(self, data_file, queries_file):
    '''
    Initialize by parsing the data table and query set files, and loading them to memory.
    '''
    super().__init__()
    log('Preparing to process query using local cache.')
    self.data_file = data_file
    dtype = {
      '#Organism/Name' : 'str',
      'TaxID' : 'int',
      'BioProject Accession' : 'str',
      'BioProject ID' : 'str', # contains '-'
      'Group' : 'str',
      'SubGroup' : 'str',
      'Size (Mb)' : 'float',
      'GC%' : 'str', # contains '-'
      'Replicons' : 'str',
      'WGS' : 'str',
      'Scaffolds' : 'str', # contains '-'
      'Genes' : 'str', # contains '-'
      'Proteins' : 'str', # contains '-'
      'Release Date' : 'str',
      'Modify Date' : 'str',
      'Status' : 'str',
      'Centert' : 'str',
      'BioSample Accession' : 'str',
      'Assembly Accession' : 'str',
      'Reference' : 'str',
      'FTP Path' : 'str',
      'Pubmed ID' : 'str', # must be parsed as CSV ints
      'Strain' : 'str',
      'Phylum' : 'str',
      'Family' : 'str',
      'Genus' : 'str',
      'Species' : 'str',
    }
    parse_dates = ['Release Date', 'Modify Date']
    log('Loading local cache in preparation for querying.', 1)
    self.table = pd.read_csv(
      data_file,
      sep='\t',
      dtype=dtype,
      parse_dates=parse_dates
    )
    self.queries = parseQueries(queries_file)
  
  # TODO: probably can refactor this into a local call to self.table.columns, and remove the convenience method.
  def columns(self):
    return self.table.columns

  def criterionTruthOf(self, criterion):
    '''
    Build truth table for which rows of the data table match the given Criterion.
    '''
    criterion_truth = None
    if (len(criterion.values) > 0):
      try:
        criterion_truth = self.table[criterion.column] == criterion.values[0]
        log('Found first value', 3)
        for additional_value in criterion.values[1:]:
          criterion_truth |= self.table[criterion.column] == additional_value
          log('Found additional value', 3)
      except KeyError:
        log(f'\n ** KeyError: The query file specifies a column name ("{criterion.column}") that does not exist in in the data file ("{self.data_file}"). Aborting. **')
        log('\n    Available columns are:\n',1)    
        {log('        ' + column, 1) for column in self.columns()}
        print()
        sys.exit()
    else:
      log(f'No values found for this criterion: {criterion}. Please fix the query file. Aborting.')
      sys.exit()
    return criterion_truth

  def queryTruthOf(self, query: Query):
    '''
    Build truth table for which rows of the data table match the given Query.
    '''
    query_truth = None
    if (len(query.criteria) > 0):
      log('Found first criterion', 3)
      query_truth = self.criterionTruthOf(query.criteria[0])
      for additional_criterion in query.criteria[1:]:
        log('Found additional criterion', 3)
        query_truth &= self.criterionTruthOf(additional_criterion)
    else:
      log(f'No criteria found for this query: {query}. Please fix the query file. Aborting.')
      sys.exit()
    return query_truth


  def select_by_queries(self):
    '''
    Select and return rows matching given query set.
    '''
    log('Applying parsed queries to data', 3)
    total_truth = None
    if (len(self.queries) > 0):
      log('Found first query', 3)
      total_truth = self.queryTruthOf(self.queries[0])
      for additional_query in self.queries[1:]:
        log('Found additional query', 3)
        total_truth |= self.queryTruthOf(additional_query)
    else:
      log('No queries found. Exiting.')
      sys.exit()
    log('Returning all matching species', 1)
    return self.table[total_truth]
        

def main():
  '''
  Run interactive command-line tool.
  '''

  log("Welcome to Proqueryote!")

  # Detect whether user has requested FNA or default FAA files,
  # and configure corresponding output strings.
  if (args.fna):
    log(f'Configured to download FNA genomes', 3)
  else:
    log(f'Configured to download FAA proteomes', 3)
  
  def fileExtension():
    if (args.fna):
      return FNA_EXTENSION
    else:
      return FAA_EXTENSION

  def dataType():
    if (args.fna):
      return 'genomes'
    else:
      return 'proteomes'

  def directoryStub():
    if (args.fna):
      return 'fna'
    else:
      return 'faa'

  # Download and install cache, if not already present.
  if not testCachePresence():
    updateCache()

  # Load new, expanded prokaryotes_taxonomic.txt data file,
  # with new added columns, to memory, for fast searching.
  data = Data(AUGMENTED, args.queries_file)
  
  # Search data table using user's query set
  selected = data.select_by_queries()
  log(f'Selected these records:\n{selected}', 3)
  count = len(selected)

  # Inform user of number of hits, and prompt to download corresponding sequence files.
  print(f'\nFound {count} species. Download available {dataType()}? [y/N]')
  while (True):
    choice = sys.stdin.readline().rstrip().lower()
    if ((choice == 'n') or (choice == '')):
      print('Aborting.')
      return
    elif (choice != 'y'):
      print('Please type "y" or "n" and then press ENTER:')
      continue
    else:
      break
  
  # Build list of sequence download URLs for search hits.
  urls = selected['FTP Path']

  # Create new folder in working directory to which to download sequences.
  # Generate unique folder name, including date-time stamp, type of
  # file being downloaded, and number of hits from this search.
  stamp = datetime.datetime.now().strftime('%Y-%m-%d-%H%M.%S')
  folder = f'proqueryote-{directoryStub()}-{count}-{stamp}'
  os.system(f'mkdir {folder}')

  # Download sequences, of chosen type, corresponding to search hits.
  os.chdir(folder)
  for url in urls:
    sys.stdout.write('.')
    sys.stdout.flush()
    prot_url = url + fileExtension()
    os.system(f'wget --quiet {prot_url}')
  # Extract sequence files from compressed archives.
  os.system('gunzip -- *')
  os.chdir('..')
  print()

# Detects whether this program is being run as a script.
# (Currently the only way to use proqueryote.)
if __name__ == "__main__":
    main()
#!/usr/bin/env python

# proqueryote
# UNH MBCBS 913 Spring 2020
# <NAME>


import os
import sys
import datetime
import argparse
import re
import pandas as pd
from typing import List
from pathlib import Path
from collections import namedtuple
Values = List[str]

'''
=========================================================
GLOBAL CONFIGURATION, INITIALIZATION, AND SUPPORT METHODS
=========================================================
'''

# CACHE FILE LOCATIONS
# Cache is stored in hidden folder in user's home folder.
# Cache contains database files used to process user queries.
# (Actual sequence files are not cached.)
HOME = str(Path.home())
CACHE = HOME + '/.proqueryote/'
PROKARYOTES = CACHE + 'prokaryotes.txt'
NODES_FILE = CACHE + 'nodes.dmp'
NAMES_FILE = CACHE + 'names.dmp'
AUGMENTED = CACHE + 'prokaryotes_taxonomic.txt'
TAXDUMP_FILENAME = 'taxdump.tar.gz'
TAXDUMP = CACHE + '/' + TAXDUMP_FILENAME
CACHE_FILES = [PROKARYOTES, NAMES_FILE, NODES_FILE, AUGMENTED]
CACHE_URLS = [
  'ftp://ftp.ncbi.nlm.nih.gov/genomes/GENOME_REPORTS/prokaryotes.txt',
  'https://ftp.ncbi.nlm.nih.gov/pub/taxonomy/' + TAXDUMP_FILENAME
]

# SPECIFY SEQUENCE FILE TYPES TO DOWNLOAD
FAA_EXTENSION = '/*protein.faa.gz'
FNA_EXTENSION = '/*_cds_from_genomic.fna.gz'

# ORDERED LISTING OF ALL RANKS USED BY NODES.DMP (FROM TAXDUMP.TAR.GZ FROM NCBI)
# This ordered list is necessary for building the taxonomy for each initial taxID listed in prokaryotes.txt
RANK_ORDER = ['superkingdom', 'kingdom', 'superphylum', 'phylum', 'subphylum', 'superclass', 'class', 'subclass', 'infraclass', 'order', 'superfamily', 'family', 'subfamily', 'genus', 'subgenus', 'species group', 'species subgroup', 'species', 'subspecies']

# RANKS FOR WHICH TO ADD COLUMNS TO PROKARYOTES TABLE
# These are the specific taxonomic ranks whose informamtion we wanted
# added as new columns in the prokaryotes.txt table.
def RANKS():
   return {'species':'','genus':'','family':'','phylum':''}

# FOR COMMAND LINE HELP INFO
epilog = """
See query file format specification in associated README.
"""

# SET UP COMMAND LINE ARGUMENTS
parser = argparse.ArgumentParser(prog="column_search", usage="column_search <queries_file>", description="Searches prokaryotes_taxanomic.txt file in working directory for all rows matching queries in given <queries_file>.", epilog=epilog, formatter_class=argparse.RawTextHelpFormatter)

parser.add_argument("queries_file", type=str)
parser.add_argument('--verbose', '-v', action='count', default=0,)
parser.add_argument('--fna', action='store_const', const=True, default=False)
args = parser.parse_args()

# PRINT A MESSAGE TO STANDARD OUTPUT PER VERBOSITY LEVEL
# Only prints messages with level equal to or less than
# verbosity level specified by user on command line.
def log(msg, level=0):
   if (args.verbose >= level): print(msg)

'''
==============
CACHE HANDLING
==============
'''

# TEST CACHE PRESENCE
# Returns False if any of the cache files, or the cache directory, is missing).
def testCachePresence():
  log('testing cache presence', 2)
  if (os.path.isdir(CACHE)):
    log('testCachePresence: cache directory present', 3)
    for each in CACHE_FILES:
      if not os.path.exists(each):
        log(f'testCachePresence: {each} is missing from cache', 3)
        return False
      else:
        log(f'testCachePresence: {each} file is present', 3)
    log('testCachePresence: cache found to be present. OK.', 3)
    return True
  else:
    log('testCachePresence: cache directory missing', 3)
    return False

# UPDATE CACHE
# Removes current cache, if present, and downloads entirely new cache.
def updateCache():
  log('Updating local cache.')
  # remove old cache, create new
  # (will need to make this more fine-grained if we ever
  # store config files in the hidden folder)
  os.system(f'rm -r {CACHE} 2> /dev/null')
  os.mkdir(CACHE)
  
  # download prokaryotes and nodes DB and extract
  working_directory = os.getcwd()
  os.chdir(CACHE)
  log('Downloading resources for local cache', 1)
  for resource in CACHE_URLS:
    log(f'{resource}', 2)
    os.system(f'wget --quiet {resource}')
  os.system(f'tar -z -xf {TAXDUMP}')
  os.chdir(working_directory)

  # created augmented version of table, with new taxonomic rank columns
  taxonomifier = Taxonomifier()
  taxonomifier.produceTaxonomic()


'''
================================================================
BUILD TAXONOMIC INFORMATION AND ADD TO NEW PROKARYOTES.TXT TABLE
================================================================
'''

# PARENT AND RANK
# A convenience type: for a given tax ID, store its rank, and the taxID of its parent.
ParentAndRank = namedtuple('ParentAndRank', 'parent rank')

# TAXONOMIFIER
# Class to organize methods and data for adding new taxonomic rank columns to prokaryotes.txt
class Taxonomifier(object):
  def __init__(self):
    log(f'Preparing to augment table with additional taxonomic rank columns', 1)
    log(f'Loading local cache', 2)
    log(f'Loading {PROKARYOTES}', 3)
    self.proks = open(PROKARYOTES, 'r').readlines()
    self.nodes = self.getNodesMap()
    self.names = self.getNamesMap()
    log(f'Read in {len(self.proks)} lines from {PROKARYOTES}', 3)

  # GET NAMES MAP
  # Load the names.dmp database to memory, for fast querying,
  # as a dictionary of taxID:name key:value pairs.
  @staticmethod
  def getNamesMap():
    log(f'Loading {NAMES_FILE}', 3)
    names = dict()
    with open(NAMES_FILE, 'r') as names_f:
        names_d = names_f.readlines()
        for record in names_d:
          cells = re.split(r'\t\|\t', record)
          taxid = cells[0]
          name = cells[1]
          name_type = cells[3]
          if (name_type.startswith('scientific name')):
              names.setdefault(int(taxid), name)
    return names

  # GET NODES MAP
  # Load the names.dmp database to memory, for fast querying,
  # as a dictionary of taxID:ParentAndRank key:value pairs.
  # (ParentAndRank is the tuple type defind above.)
  @staticmethod
  def getNodesMap():
    log(f'Loading {NODES_FILE}', 3)
    nodes = dict()
    with open(NODES_FILE, 'r') as nodes_f:
        nodes_d = nodes_f.readlines()
        for record in nodes_d:
          cells = re.split(r'\t\|\t', record)
          this_taxid = cells[0]
          this_parent_taxid = cells[1]
          this_rank = cells[2]
          parent_rank = ParentAndRank(int(this_parent_taxid), this_rank)
          nodes.setdefault(int(this_taxid),parent_rank)
    return nodes

  def getName(self, taxid):
    """
    Retrieves scientific name corresponding to given taxid.
    """
    return self.names[int(taxid)]

  def getParent(self, taxid):
    """
    Retrieves taxid of parent of given taxid.
    """
    return self.nodes[int(taxid)].parent

  def getRank(self, taxid):
    """
    Retrieves rank of given taxid.
    """
    return self.nodes[int(taxid)].rank

  @staticmethod
  def getTaxID(proks_record):
    """
    Parse TaxID field from raw string of given row from prokaryotes.txt
    """
    return proks_record.split('\t')[1]

  def getOneTaxonomy(self, taxid):
    """
    Searches nodes.dmp for desired ranks, pulling corresponding names from names.dmp.
    This method builds the string that can be appended to a given row from prokaryotes.txt.

    :return: A TSV string: "<species>\t<genus>\t<family>\t<phylum>"
    """
    ranks = RANKS()
    searchTaxID = str(taxid)
    log(f'getOneTaxonomy: taxid {taxid}', 4)
    # continue searching until we've found an entry for each desired taxonomic rank
    while ('' in ranks.values()):
        log(f'getOneTaxonomy: searching for {searchTaxID}', 4)
        if (searchTaxID == '1'):
          log(f'getOneTaxonomy: arrived at root, curtailing search with extant results: {ranks}', 4)
          break
        try:
          this_parent_taxid = self.getParent(searchTaxID)
        except KeyError:
          log(f'getOneTaxonomy: caught KeyError raised by getParent: TaxID {searchTaxID} is not listed in {args.nodes_file}.', 4)
          break
        this_rank = self.getRank(searchTaxID)
        log(f'getOneTaxonomy: {searchTaxID} is a {this_rank} with parent {this_parent_taxid}', 4)
        # TODO: may want to alter this to confirm that ranks[this_rank] == '',
        # to prevent over-writing of multiple entries exist in nodes.dmp, not all correct?:
        if (this_rank in ranks):
          name = self.getName(searchTaxID)
          ranks[this_rank] = name
          log(f'getOneTaxonomy: added {searchTaxID} : {name} as {this_rank}. we have thus far found: {ranks}', 4)
        searchTaxID = str(this_parent_taxid)
    log(f'getOneTaxonomy: found {ranks}', 4)
    return(ranks)

  def produceTaxonomic(self):
    """
    Create a new prokaryotes_taxonomic.txt table file, with additional taxonomic rank columns.
    """
    log('Building new table', 1)
    # costruct new header row
    old_header = self.proks[0].rstrip()
    new_cols = RANKS().keys()
    log('old header:\n' + old_header, 3)
    log('new columns:\n' + str(new_cols), 4)
    new_header = old_header
    for rank in RANK_ORDER:
      if (rank in new_cols):
        new_header += '\t' + rank.title()
    new_header += '\n'
    log('new header:\n' + new_header, 3)

    # start writing to new file
    with open(AUGMENTED, 'w') as out_f:
      # write new header row
      log(f'About to write new header to new, augmented table:\n{new_header}', 4)
      out_f.write(new_header)
      
      # construct new entries
      for record in self.proks[1:]:
        log('old row:\n' + record, 4)
        taxid = self.getTaxID(record)
        taxonomy = self.getOneTaxonomy(taxid)
        new_record = record.rstrip()
        for rank in RANK_ORDER:
          if (rank in taxonomy.keys()):
              new_record += '\t' + taxonomy[rank]
        new_record += '\n'
        log('new record:\n' + new_record, 4)
        out_f.write(new_record)
    log(f'New table with additional taxonomic columns:\n\t{AUGMENTED}', 2)





'''
=======================================================================================================
PARSE AND PROCESS QUERIES, NOTIFY USER OF NUMBER OF RESULTS, AND PROMPT TO DOWNLOAD SELECTED SEQUENCES.
=======================================================================================================
'''

class Criterion(object):
  '''
  Specifies acceptable values for a given column, for a given query.
  '''
  def __init__(self, column: str, values: Values):
    self.column = column
    self.values = values
  def __repr__(self):
    return f'{self.column}: {self.values}'

class Query(object):
  '''
  Specifies all criteria (in form of Criterion objects) for a given query.
  '''
  def __init__(self):
    self.criteria = list()
  
  def addCriterion(self, criterion: Criterion):
    self.criteria.append(criterion)
  
  def __repr__(self):
    return 'Query:\n' + str(self.criteria)

def parseQueries(queries_file):
  '''
  Parse a query set file, in preparation for processing/executing the queries contained therein.
  '''
  log(f'Parsing query file: {queries_file}', 1)
  queries = list()
  with open(queries_file, 'r') as queries_f:
    for line in queries_f.readlines():
      log("  starting loop", 3)
      line = line.rstrip()
      # skip blank and comment lines
      if (line == '' or line[0] == '#'):
        log("  found empty line or comment; skipping", 3)
        continue
      # start a new query
      elif (line == 'query'):
        log("  found new query", 3)
        queries.append(Query())
      # append criterion to current query
      else:
        log("  adding to criterion to current query", 3)
        delimited = re.split(' ', line)
        queries[len(queries) - 1].addCriterion(Criterion(delimited[0],delimited[1:]))        
  log(f'parsed queries:\n{queries}', 3)
  return queries


class Data(object):
  '''
  Class organizing the data table to be queried, the parsed
  queries, and methods for processing queries against the data table.
  '''
  def __init__(self, data_file, queries_file):
    '''
    Initialize by parsing the data table and query set files, and loading them to memory.
    '''
    super().__init__()
    log('Preparing to process query using local cache.')
    self.data_file = data_file
    dtype = {
      '#Organism/Name' : 'str',
      'TaxID' : 'int',
      'BioProject Accession' : 'str',
      'BioProject ID' : 'str', # contains '-'
      'Group' : 'str',
      'SubGroup' : 'str',
      'Size (Mb)' : 'float',
      'GC%' : 'str', # contains '-'
      'Replicons' : 'str',
      'WGS' : 'str',
      'Scaffolds' : 'str', # contains '-'
      'Genes' : 'str', # contains '-'
      'Proteins' : 'str', # contains '-'
      'Release Date' : 'str',
      'Modify Date' : 'str',
      'Status' : 'str',
      'Centert' : 'str',
      'BioSample Accession' : 'str',
      'Assembly Accession' : 'str',
      'Reference' : 'str',
      'FTP Path' : 'str',
      'Pubmed ID' : 'str', # must be parsed as CSV ints
      'Strain' : 'str',
      'Phylum' : 'str',
      'Family' : 'str',
      'Genus' : 'str',
      'Species' : 'str',
    }
    parse_dates = ['Release Date', 'Modify Date']
    log('Loading local cache in preparation for querying.', 1)
    self.table = pd.read_csv(
      data_file,
      sep='\t',
      dtype=dtype,
      parse_dates=parse_dates
    )
    self.queries = parseQueries(queries_file)
  
  # TODO: probably can refactor this into a local call to self.table.columns, and remove the convenience method.
  def columns(self):
    return self.table.columns

  def criterionTruthOf(self, criterion):
    '''
    Build truth table for which rows of the data table match the given Criterion.
    '''
    criterion_truth = None
    if (len(criterion.values) > 0):
      try:
        criterion_truth = self.table[criterion.column] == criterion.values[0]
        log('Found first value', 3)
        for additional_value in criterion.values[1:]:
          criterion_truth |= self.table[criterion.column] == additional_value
          log('Found additional value', 3)
      except KeyError:
        log(f'\n ** KeyError: The query file specifies a column name ("{criterion.column}") that does not exist in in the data file ("{self.data_file}"). Aborting. **')
        log('\n    Available columns are:\n',1)    
        {log('        ' + column, 1) for column in self.columns()}
        print()
        sys.exit()
    else:
      log(f'No values found for this criterion: {criterion}. Please fix the query file. Aborting.')
      sys.exit()
    return criterion_truth

  def queryTruthOf(self, query: Query):
    '''
    Build truth table for which rows of the data table match the given Query.
    '''
    query_truth = None
    if (len(query.criteria) > 0):
      log('Found first criterion', 3)
      query_truth = self.criterionTruthOf(query.criteria[0])
      for additional_criterion in query.criteria[1:]:
        log('Found additional criterion', 3)
        query_truth &= self.criterionTruthOf(additional_criterion)
    else:
      log(f'No criteria found for this query: {query}. Please fix the query file. Aborting.')
      sys.exit()
    return query_truth


  def select_by_queries(self):
    '''
    Select and return rows matching given query set.
    '''
    log('Applying parsed queries to data', 3)
    total_truth = None
    if (len(self.queries) > 0):
      log('Found first query', 3)
      total_truth = self.queryTruthOf(self.queries[0])
      for additional_query in self.queries[1:]:
        log('Found additional query', 3)
        total_truth |= self.queryTruthOf(additional_query)
    else:
      log('No queries found. Exiting.')
      sys.exit()
    log('Returning all matching species', 1)
    return self.table[total_truth]
        

def main():
  '''
  Run interactive command-line tool.
  '''

  log("Welcome to Proqueryote!")

  # Detect whether user has requested FNA or default FAA files,
  # and configure corresponding output strings.
  if (args.fna):
    log(f'Configured to download FNA genomes', 3)
  else:
    log(f'Configured to download FAA proteomes', 3)
  
  def fileExtension():
    if (args.fna):
      return FNA_EXTENSION
    else:
      return FAA_EXTENSION

  def dataType():
    if (args.fna):
      return 'genomes'
    else:
      return 'proteomes'

  def directoryStub():
    if (args.fna):
      return 'fna'
    else:
      return 'faa'

  # Download and install cache, if not already present.
  if not testCachePresence():
    updateCache()

  # Load new, expanded prokaryotes_taxonomic.txt data file,
  # with new added columns, to memory, for fast searching.
  data = Data(AUGMENTED, args.queries_file)
  
  # Search data table using user's query set
  selected = data.select_by_queries()
  log(f'Selected these records:\n{selected}', 3)
  count = len(selected)

  # Inform user of number of hits, and prompt to download corresponding sequence files.
  print(f'\nFound {count} species. Download available {dataType()}? [y/N]')
  while (True):
    choice = sys.stdin.readline().rstrip().lower()
    if ((choice == 'n') or (choice == '')):
      print('Aborting.')
      return
    elif (choice != 'y'):
      print('Please type "y" or "n" and then press ENTER:')
      continue
    else:
      break
  
  # Build list of sequence download URLs for search hits.
  urls = selected['FTP Path']

  # Create new folder in working directory to which to download sequences.
  # Generate unique folder name, including date-time stamp, type of
  # file being downloaded, and number of hits from this search.
  stamp = datetime.datetime.now().strftime('%Y-%m-%d-%H%M.%S')
  folder = f'proqueryote-{directoryStub()}-{count}-{stamp}'
  os.system(f'mkdir {folder}')

  # Download sequences, of chosen type, corresponding to search hits.
  os.chdir(folder)
  for url in urls:
    sys.stdout.write('.')
    sys.stdout.flush()
    prot_url = url + fileExtension()
    os.system(f'wget --quiet {prot_url}')
  # Extract sequence files from compressed archives.
  os.system('gunzip -- *')
  os.chdir('..')
  print()

# Detects whether this program is being run as a script.
# (Currently the only way to use proqueryote.)
if __name__ == "__main__":
    main()
#!/usr/bin/env python # proqueryote # UNH MBCBS 913 Spring 2020 # <NAME> ========================================================= GLOBAL CONFIGURATION, INITIALIZATION, AND SUPPORT METHODS ========================================================= # CACHE FILE LOCATIONS # Cache is stored in hidden folder in user's home folder. # Cache contains database files used to process user queries. # (Actual sequence files are not cached.) # SPECIFY SEQUENCE FILE TYPES TO DOWNLOAD # ORDERED LISTING OF ALL RANKS USED BY NODES.DMP (FROM TAXDUMP.TAR.GZ FROM NCBI) # This ordered list is necessary for building the taxonomy for each initial taxID listed in prokaryotes.txt # RANKS FOR WHICH TO ADD COLUMNS TO PROKARYOTES TABLE # These are the specific taxonomic ranks whose informamtion we wanted # added as new columns in the prokaryotes.txt table. # FOR COMMAND LINE HELP INFO See query file format specification in associated README. # SET UP COMMAND LINE ARGUMENTS # PRINT A MESSAGE TO STANDARD OUTPUT PER VERBOSITY LEVEL # Only prints messages with level equal to or less than # verbosity level specified by user on command line. ============== CACHE HANDLING ============== # TEST CACHE PRESENCE # Returns False if any of the cache files, or the cache directory, is missing). # UPDATE CACHE # Removes current cache, if present, and downloads entirely new cache. # remove old cache, create new # (will need to make this more fine-grained if we ever # store config files in the hidden folder) # download prokaryotes and nodes DB and extract # created augmented version of table, with new taxonomic rank columns ================================================================ BUILD TAXONOMIC INFORMATION AND ADD TO NEW PROKARYOTES.TXT TABLE ================================================================ # PARENT AND RANK # A convenience type: for a given tax ID, store its rank, and the taxID of its parent. # TAXONOMIFIER # Class to organize methods and data for adding new taxonomic rank columns to prokaryotes.txt # GET NAMES MAP # Load the names.dmp database to memory, for fast querying, # as a dictionary of taxID:name key:value pairs. # GET NODES MAP # Load the names.dmp database to memory, for fast querying, # as a dictionary of taxID:ParentAndRank key:value pairs. # (ParentAndRank is the tuple type defind above.) Retrieves scientific name corresponding to given taxid. Retrieves taxid of parent of given taxid. Retrieves rank of given taxid. Parse TaxID field from raw string of given row from prokaryotes.txt Searches nodes.dmp for desired ranks, pulling corresponding names from names.dmp.     This method builds the string that can be appended to a given row from prokaryotes.txt.      :return: A TSV string: "<species>\t<genus>\t<family>\t<phylum>" # continue searching until we've found an entry for each desired taxonomic rank # TODO: may want to alter this to confirm that ranks[this_rank] == '', # to prevent over-writing of multiple entries exist in nodes.dmp, not all correct?: Create a new prokaryotes_taxonomic.txt table file, with additional taxonomic rank columns. # costruct new header row # start writing to new file # write new header row # construct new entries ======================================================================================================= PARSE AND PROCESS QUERIES, NOTIFY USER OF NUMBER OF RESULTS, AND PROMPT TO DOWNLOAD SELECTED SEQUENCES. ======================================================================================================= Specifies acceptable values for a given column, for a given query. Specifies all criteria (in form of Criterion objects) for a given query. Parse a query set file, in preparation for processing/executing the queries contained therein. # skip blank and comment lines # start a new query # append criterion to current query Class organizing the data table to be queried, the parsed   queries, and methods for processing queries against the data table. Initialize by parsing the data table and query set files, and loading them to memory. # contains '-' # contains '-' # contains '-' # contains '-' # contains '-' # must be parsed as CSV ints # TODO: probably can refactor this into a local call to self.table.columns, and remove the convenience method. Build truth table for which rows of the data table match the given Criterion. Build truth table for which rows of the data table match the given Query. Select and return rows matching given query set. Run interactive command-line tool. # Detect whether user has requested FNA or default FAA files, # and configure corresponding output strings. # Download and install cache, if not already present. # Load new, expanded prokaryotes_taxonomic.txt data file, # with new added columns, to memory, for fast searching. # Search data table using user's query set # Inform user of number of hits, and prompt to download corresponding sequence files. # Build list of sequence download URLs for search hits. # Create new folder in working directory to which to download sequences. # Generate unique folder name, including date-time stamp, type of # file being downloaded, and number of hits from this search. # Download sequences, of chosen type, corresponding to search hits. # Extract sequence files from compressed archives. # Detects whether this program is being run as a script. # (Currently the only way to use proqueryote.)