diff --git a/protmapper/resources.py b/protmapper/resources.py index b7f0501..4b36306 100644 --- a/protmapper/resources.py +++ b/protmapper/resources.py @@ -192,6 +192,158 @@ def process_uniprot_line(line, base_columns, processed_columns, return '\t'.join(new_line) +def _validate_taxonomy_id(taxonomy_id): + """Return an NCBI taxonomy ID as a string""" + tid = str(taxonomy_id).strip() + if not tid.isdigit(): + raise ValueError('Invalid NCBI taxonomy ID: %s' % taxonomy_id) + return tid + + +def build_organism_query(taxonomy_ids, include_unreviewed=True): + """Return a UniProt query string for entries from the given organisms. + + Parameters + ---------- + taxonomy_ids : list[str] + A list of NCBI taxonomy IDs + include_unreviewed : bool + If True, the query matches both reviewed and unreviewed entries + Returns + ------- + str + A UniProt query string. + """ + if isinstance(taxonomy_ids, str): + raise ValueError('Expected a list of taxonomy IDs but got a single ' + 'string: %s' % taxonomy_ids) + tids = [] + for taxonomy_id in taxonomy_ids: + tid = _validate_taxonomy_id(taxonomy_id) + if tid not in tids: + tids.append(tid) + if not tids: + raise ValueError('At least one taxonomy ID is required.') + query = '(%s)' % ' OR '.join('taxonomy_id:%s' % tid for tid in tids) + reviewed = ('(reviewed:true OR reviewed:false)' if include_unreviewed + else 'reviewed:true') + return '%s AND %s' % (reviewed, query) + + +def _iter_uniprot_search_lines(query, columns, page_size=500): + """Yield lines of a UniProt search result in TSV format. + + The stream endpoint used elsewhere in this module returns an entire + result set in a single response, which doesn't scale once unreviewed + entries are included. The search endpoint used here instead pages + through results. + + Parameters + ---------- + query : str + A UniProt query string. + columns : list[str] + The UniProt fields to download. + page_size : int + The number of entries to request per page, at most 500. + + Yields + ------ + str + The header line of the TSV table, followed by one line per entry. + """ + url = 'https://rest.uniprot.org/uniprotkb/search' + params = { + 'format': 'tsv', + 'query': query, + 'size': page_size, + 'fields': ','.join(columns), + 'sort': 'accession asc', + } + header = None + num_rows = 0 + total = None + while url: + res = requests.get(url, params=params) + # Pages after the first are fetched from the cursor URL given in the + # Link header, which already carries all the query parameters. + params = None + if res.status_code != 200: + raise RuntimeError('Failed to download "%s": status %s' % + (res.url, res.status_code)) + if total is None: + total = res.headers.get('x-total-results') + logger.info('Downloading %s UniProt entries' % + (total if total else 'an unknown number of')) + lines = res.text.splitlines() + if lines: + # Each page repeats the header which we only emit once + if header is None: + header = lines[0] + yield header + for line in lines[1:]: + num_rows += 1 + yield line + url = res.links.get('next', {}).get('url') + if not num_rows: + # Note that this only catches a query that matched nothing at all. + # Since the taxa of a query are OR-ed together, a single taxon that + # UniProt no longer annotates any entry against contributes nothing + # without being noticed here. Checking that requires querying each + # taxon on its own. + logger.warning('No UniProt entries found for query: %s' % query) + else: + logger.info('Downloaded %d UniProt entries' % num_rows) + + +def download_uniprot_entries_for_organisms(out_file, taxonomy_ids, columns, + include_unreviewed=True, + cached=True): + """Download UniProt entries for a given set of organisms. + + Parameters + ---------- + out_file : str + Path of the compressed TSV file to write the entries into. + taxonomy_ids : list[str] + A list of NCBI taxonomy IDs. Entries from + any child taxon of the given taxa, such as specific strains of a + virus species, are included too. + columns : list[str] + The UniProt fields to download, e.g., ['accession', 'protein_name']. + These are the names of the fields rather than the labels that UniProt + puts in the header of the resulting table, so 'organism_id' here + gives a column headed "Organism (ID)". The valid names are listed at + https://www.uniprot.org/help/return_fields + include_unreviewed : bool + If True, both reviewed and unreviewed entries + are downloaded, otherwise only reviewed ones. + cached : bool + If True and out_file already exists, nothing is downloaded. + Note that this differs from the other download functions + in this module, whose cache is a pre-processed file on S3. S + So the cache is just the given file. Since the + contents depend on the requested organisms, pass False when + the requested organisms change. + """ + if cached and os.path.exists(out_file): + logger.info('Using the UniProt entries already at %s' % out_file) + return + query = build_organism_query(taxonomy_ids, + include_unreviewed=include_unreviewed) + logger.info('Downloading UniProt entries for query: %s' % query) + # We write to a temporary file and only move it into place once the + # download has finished. Otherwise an interrupted download leaves a + # partial file at out_file, which a later call with cached=True has no + # way of telling apart from a complete one. + temp_file = out_file + '.tmp' + with gzip.open(temp_file, 'wt', encoding='utf-8') as fh: + for line in _iter_uniprot_search_lines(query, columns): + fh.write(line + '\n') + os.replace(temp_file, out_file) + logger.info('Saved UniProt entries into %s' % out_file) + + def parse_uniprot_synonyms(synonyms_str): synonyms_str = re.sub(r'\[Includes: ([^]])+\]', '', synonyms_str).strip()