protein-quest 0.6.0__py3-none-any.whl → 0.8.0__py3-none-any.whl

protein_quest/mcp_server.py

@@ -45,10 +45,14 @@ from protein_quest.alphafold.fetch import AlphaFoldEntry, DownloadableFormat
 from protein_quest.alphafold.fetch import fetch_many as alphafold_fetch
 from protein_quest.emdb import fetch as emdb_fetch
 from protein_quest.go import search_gene_ontology_term
-from protein_quest.io import convert_to_cif_file, glob_structure_files
+from protein_quest.io import convert_to_cif_file, glob_structure_files, read_structure
 from protein_quest.pdbe.fetch import fetch as pdbe_fetch
 from protein_quest.ss import filter_file_on_secondary_structure
-from protein_quest.structure import nr_residues_in_chain, write_single_chain_structure_file
+from protein_quest.structure import (
+    nr_residues_in_chain,
+    structure2uniprot_accessions,
+    write_single_chain_structure_file,
+)
 from protein_quest.taxonomy import search_taxon
 from protein_quest.uniprot import (
     PdbResult,
@@ -129,7 +133,7 @@ def extract_single_chain_from_structure(
 
 @mcp.tool
 def list_structure_files(path: Path) -> list[Path]:
-    """List structure files (.pdb, .pdb.gz, .cif, .cif.gz) in the specified directory."""
+    """List structure files (.pdb, .pdb.gz, .cif, .cif.gz, .bcif) in the specified directory."""
     return list(glob_structure_files(path))
 
 
@@ -150,7 +154,7 @@ def search_alphafolds(
     Field(description="Set of uniprot accessions which have an AlphaFold entry"),
 ]:
     """Search for AlphaFold entries in UniProtKB accessions."""
-    # each uniprot accesion can have one or more AlphaFold IDs
+    # each uniprot accession can have one or more AlphaFold IDs
     # an AlphaFold ID is the same as the uniprot accession
     # so we return a subset of uniprot_accs
     results = search4af(uniprot_accs, limit)
@@ -163,7 +167,7 @@ mcp.tool(search4macromolecular_complexes, name="search_macromolecular_complexes"
 
 @mcp.tool
 def fetch_alphafold_structures(uniprot_accs: set[str], save_dir: Path) -> list[AlphaFoldEntry]:
-    """Fetch the AlphaFold summary and mmcif file for given UniProt accessions.
+    """Fetch the AlphaFold mmCIF file for given UniProt accessions.
 
     Args:
         uniprot_accs: A set of UniProt accessions.
@@ -172,8 +176,8 @@ def fetch_alphafold_structures(uniprot_accs: set[str], save_dir: Path) -> list[A
     Returns:
         A list of AlphaFold entries.
     """
-    what: set[DownloadableFormat] = {"summary", "cif"}
-    return alphafold_fetch(uniprot_accs, save_dir, what)
+    formats: set[DownloadableFormat] = {"cif"}
+    return alphafold_fetch(uniprot_accs, save_dir, formats)
 
 
 @mcp.tool
@@ -203,6 +207,13 @@ mcp.tool(filter_file_on_secondary_structure)
 mcp.tool(convert_to_cif_file)
 
 
+@mcp.tool
+def uniprot_accessions_of_structure_file(file: Path) -> set[str]:
+    """Extract UniProt accessions from structure file."""
+    structure = read_structure(file)
+    return structure2uniprot_accessions(structure)
+
+
 @mcp.prompt
 def candidate_structures(
     species: str = "Human",

protein_quest/structure.py

@@ -206,3 +206,27 @@ def write_single_chain_structure_file(
     write_structure(structure, output_file)
 
     return output_file
+
+
+def structure2uniprot_accessions(structure: gemmi.Structure) -> set[str]:
+    """Extract UniProt accessions from a gemmi Structure object.
+
+    Logs a warning and returns an empty set if no accessions are found in structure.
+
+    Args:
+        structure: The gemmi Structure object to extract UniProt accessions from.
+
+    Returns:
+        A set of UniProt accessions found in the structure.
+    """
+    block = structure.make_mmcif_block(gemmi.MmcifOutputGroups(False, struct_ref=True))
+    struct_ref = block.get_mmcif_category("_struct_ref.")
+    uniprot_accessions: set[str] = set()
+    for i, db_name in enumerate(struct_ref["db_name"]):
+        if db_name != "UNP":
+            continue
+        pdbx_db_accession = struct_ref["pdbx_db_accession"][i]
+        uniprot_accessions.add(pdbx_db_accession)
+    if not uniprot_accessions:
+        logger.warning("No UniProt accessions found in structure %s", structure.name)
+    return uniprot_accessions

protein_quest/uniprot.py

@@ -1,8 +1,9 @@
 """Module for searching UniProtKB using SPARQL."""
 
 import logging
-from collections.abc import Collection, Iterable
+from collections.abc import Collection, Generator, Iterable
 from dataclasses import dataclass
+from functools import cached_property
 from itertools import batched
 from textwrap import dedent
 
@@ -24,6 +25,8 @@ class Query:
             (e.g., ["GO:0005634"]) or a collection of GO terms (e.g., ["GO:0005634", "GO:0005737"]).
         molecular_function_go: Molecular function in GO format. Can be a single GO term
             (e.g., ["GO:0003674"]) or a collection of GO terms (e.g., ["GO:0003674", "GO:0008150"]).
+        min_sequence_length: Minimum length of the canonical sequence.
+        max_sequence_length: Maximum length of the canonical sequence.
     """
 
     # TODO make taxon_id an int
@@ -32,6 +35,8 @@ class Query:
     subcellular_location_uniprot: str | None = None
     subcellular_location_go: list[str] | None = None
     molecular_function_go: list[str] | None = None
+    min_sequence_length: int | None = None
+    max_sequence_length: int | None = None
 
 
 def _first_chain_from_uniprot_chains(uniprot_chains: str) -> str:
@@ -39,16 +44,17 @@ def _first_chain_from_uniprot_chains(uniprot_chains: str) -> str:
 
     The UniProt chains string is formatted (with EBNF notation) as follows:
 
-        chain_group(=range)?(,chain_group(=range)?)*
+        chain_group=range(,chain_group=range)*
 
     where:
         chain_group := chain_id(/chain_id)*
-        chain_id    := [A-Za-z]+
+        chain_id    := [A-Za-z0-9]+
         range       := start-end
         start, end  := integer
 
     Args:
         uniprot_chains: A string representing UniProt chains, For example "B/D=1-81".
+
     Returns:
         The first chain identifier from the UniProt chain string. For example "B".
     """
@@ -66,6 +72,35 @@ def _first_chain_from_uniprot_chains(uniprot_chains: str) -> str:
     return chain
 
 
+def _chain_length_from_uniprot_chains(uniprot_chains: str) -> int:
+    """Calculates the total length of chain from a UniProt chains string.
+
+    See `_first_chain_from_uniprot_chains` for the format of the UniProt chains string.
+
+    Args:
+        uniprot_chains: A string representing UniProt chains, For example "B/D=1-81".
+
+    Returns:
+        The length of the chain in the UniProt chain string. For example 81 for "B/D=1-81".
+    """
+    total_length = 0
+    chains = uniprot_chains.split(",")
+    for chain in chains:
+        _, rangestr = chain.split("=")
+        start, stop = rangestr.split("-")
+        # Residue positions are 1-based so + 1
+        total_length += int(stop) - int(start) + 1
+    return total_length
+
+
+class PdbChainLengthError(ValueError):
+    """Raised when a UniProt chain description does not yield a chain length."""
+
+    def __init__(self, pdb_id: str, uniprot_chains: str):
+        msg = f"Could not determine chain length of '{pdb_id}' from '{uniprot_chains}'"
+        super().__init__(msg)
+
+
 @dataclass(frozen=True)
 class PdbResult:
     """Result of a PDB search in UniProtKB.
@@ -82,11 +117,78 @@ class PdbResult:
     uniprot_chains: str
     resolution: str | None = None
 
-    @property
+    @cached_property
     def chain(self) -> str:
         """The first chain from the UniProt chains aka self.uniprot_chains."""
         return _first_chain_from_uniprot_chains(self.uniprot_chains)
 
+    @cached_property
+    def chain_length(self) -> int:
+        """The length of the chain from the UniProt chains aka self.uniprot_chains."""
+        try:
+            return _chain_length_from_uniprot_chains(self.uniprot_chains)
+        except ValueError as e:
+            raise PdbChainLengthError(self.id, self.uniprot_chains) from e
+
+
+type PdbResults = dict[str, set[PdbResult]]
+"""Dictionary with uniprot accessions as keys and sets of PDB results as values."""
+
+
+def filter_pdb_results_on_chain_length(
+    pdb_results: PdbResults,
+    min_residues: int | None,
+    max_residues: int | None,
+    keep_invalid: bool = False,
+) -> PdbResults:
+    """Filter PDB results based on chain length.
+
+    Args:
+        pdb_results: Dictionary with protein IDs as keys and sets of PDB results as values.
+        min_residues: Minimum number of residues required in the chain mapped to the UniProt accession.
+            If None, no minimum is applied.
+        max_residues: Maximum number of residues allowed in chain mapped to the UniProt accession.
+            If None, no maximum is applied.
+        keep_invalid: If True, PDB results with invalid chain length (could not be determined) are kept.
+            If False, PDB results with invalid chain length are filtered out.
+            Warnings are logged when length can not be determined.
+
+    Returns:
+        Filtered dictionary with protein IDs as keys and sets of PDB results as values.
+    """
+    if min_residues is None and max_residues is None:
+        # No filtering needed
+        return pdb_results
+    if min_residues is not None and max_residues is not None and max_residues <= min_residues:
+        msg = f"Maximum number of residues ({max_residues}) must be > minimum number of residues ({min_residues})"
+        raise ValueError(msg)
+    results: PdbResults = {}
+    for uniprot_accession, pdb_entries in pdb_results.items():
+        filtered_pdb_entries = set()
+        for pdb_entry in pdb_entries:
+            try:
+                if (min_residues is None or pdb_entry.chain_length >= min_residues) and (
+                    max_residues is None or pdb_entry.chain_length <= max_residues
+                ):
+                    filtered_pdb_entries.add(pdb_entry)
+            except PdbChainLengthError:
+                if keep_invalid:
+                    logger.warning(
+                        f"Could not determine chain length of '{pdb_entry.id}' from '{pdb_entry.uniprot_chains}' "
+                        f"belonging to uniprot accession '{uniprot_accession}', "
+                        "for completeness not filtering it out"
+                    )
+                    filtered_pdb_entries.add(pdb_entry)
+                else:
+                    logger.warning(
+                        f"Filtering out PDB entry '{pdb_entry.id}' belonging to uniprot accession "
+                        f"'{uniprot_accession}' due to invalid chain length from '{pdb_entry.uniprot_chains}'"
+                    )
+        if filtered_pdb_entries:
+            # Only include uniprot_accession if there are any pdb entries left after filtering
+            results[uniprot_accession] = filtered_pdb_entries
+    return results
+
 
 def _query2dynamic_sparql_triples(query: Query):
     parts: list[str] = []
@@ -110,6 +212,13 @@ def _query2dynamic_sparql_triples(query: Query):
         molecular_function_filter = _create_go_filter(go_terms, "Molecular function")
         parts.append(molecular_function_filter)
 
+    if query.min_sequence_length is not None or query.max_sequence_length is not None:
+        length_filter = _build_sparql_query_sequence_length_filter(
+            min_length=query.min_sequence_length,
+            max_length=query.max_sequence_length,
+        )
+        parts.append(length_filter)
+
     return "\n".join(parts)
 
 
@@ -237,6 +346,57 @@ def _build_sparql_query_uniprot(query: Query, limit=10_000) -> str:
     return _build_sparql_generic_query(select_clause, dedent(where_clause), limit)
 
 
+def _build_sparql_query_sequence_length_filter(min_length: int | None = None, max_length: int | None = None) -> str:
+    """Builds a SPARQL filter for sequence length.
+
+    See 107_uniprot_sequences_and_mark_which_is_cannonical_for_human
+    on https://sparql.uniprot.org/.well-known/sparql-examples/ for similar query.
+
+    Args:
+        min_length: Minimum sequence length. If None, no minimum is applied.
+        max_length: Maximum sequence length. If None, no maximum is applied.
+    """
+    if min_length is None and max_length is None:
+        return ""
+    # An uniprot entry can have multiple isoforms,
+    # we want to check the length of the canonical isoform
+    # We do this by selecting the isoform that is not based on another isoform
+    # and excluding isoforms from other uniprot entries.
+    # For example for http://purl.uniprot.org/uniprot/P42284:
+    # - http://purl.uniprot.org/isoforms/P42284-2 is ok
+    # - http://purl.uniprot.org/isoforms/P42284-1 is not ok, because it is based on P42284-2
+    # - http://purl.uniprot.org/isoforms/Q7KQZ4-1 is not ok, because it is from another uniprot entry
+    header = dedent("""\
+        ?protein up:sequence ?isoform .
+        ?isoform a up:Simple_Sequence .
+        BIND (IRI(STRBEFORE(REPLACE(
+            STR(?isoform), "http://purl.uniprot.org/isoforms/", "http://purl.uniprot.org/uniprot/"
+        ), "-")) AS ?ac_of_isoform)
+        FILTER (?protein = ?ac_of_isoform)
+        ?isoform rdf:value ?sequence .
+        BIND (STRLEN(?sequence) AS ?seq_length)
+    """)
+    if min_length is not None and max_length is not None:
+        if max_length <= min_length:
+            msg = f"Maximum sequence length ({max_length}) must be greater than minimum sequence length ({min_length})"
+            raise ValueError(msg)
+        return dedent(f"""\
+            {header}
+            FILTER (?seq_length >= {min_length} && ?seq_length <= {max_length})
+        """)
+    if min_length is not None:
+        return dedent(f"""\
+            {header}
+            FILTER (?seq_length >= {min_length})
+        """)
+    if max_length is not None:
+        return dedent(f"""\
+            {header}
+            FILTER (?seq_length <= {max_length})
+        """)
+    return ""
+
+
 def _build_sparql_query_pdb(uniprot_accs: Iterable[str], limit=10_000) -> str:
     # For http://purl.uniprot.org/uniprot/O00268 + http://rdf.wwpdb.org/pdb/1H3O
     # the chainSequenceMapping are
@@ -248,7 +408,7 @@ def _build_sparql_query_pdb(uniprot_accs: Iterable[str], limit=10_000) -> str:
     # http://purl.uniprot.org/isoforms/O00255-2#PDB_3U84_tt2tt459
     # To get the the chain belonging to the uniprot/pdb pair we need to
     # do some string filtering.
-    # Also there can be multiple cnhins for the same uniprot/pdb pair, so we need to
+    # Also there can be multiple chains for the same uniprot/pdb pair, so we need to
     # do a group by and concat
 
     select_clause = dedent("""\
@@ -274,7 +434,12 @@ def _build_sparql_query_pdb(uniprot_accs: Iterable[str], limit=10_000) -> str:
     )
 
 
-def _build_sparql_query_af(uniprot_accs: Iterable[str], limit=10_000) -> str:
+def _build_sparql_query_af(
+    uniprot_accs: Iterable[str],
+    min_sequence_length: int | None = None,
+    max_sequence_length: int | None = None,
+    limit=10_000,
+) -> str:
     select_clause = "?protein ?af_db"
     where_clause = dedent("""
         # --- Protein Selection ---
@@ -284,6 +449,12 @@ def _build_sparql_query_af(uniprot_accs: Iterable[str], limit=10_000) -> str:
         ?protein rdfs:seeAlso ?af_db .
         ?af_db up:database <http://purl.uniprot.org/database/AlphaFoldDB> .
     """)
+    if min_sequence_length is not None or max_sequence_length is not None:
+        length_filter = _build_sparql_query_sequence_length_filter(
+            min_length=min_sequence_length,
+            max_length=max_sequence_length,
+        )
+        where_clause += "\n" + length_filter
     return _build_sparql_generic_by_uniprot_accessions_query(uniprot_accs, select_clause, dedent(where_clause), limit)
 
 
@@ -337,8 +508,8 @@ def _execute_sparql_search(
     return bindings
 
 
-def _flatten_results_pdb(rawresults: Iterable) -> dict[str, set[PdbResult]]:
-    pdb_entries: dict[str, set[PdbResult]] = {}
+def _flatten_results_pdb(rawresults: Iterable) -> PdbResults:
+    pdb_entries: PdbResults = {}
     for result in rawresults:
         protein = result["protein"]["value"].split("/")[-1]
         if "pdb_db" not in result:  # Should not happen with build_sparql_query_pdb
@@ -424,7 +595,7 @@ def search4uniprot(query: Query, limit: int = 10_000, timeout: int = 1_800) -> s
 
 def search4pdb(
     uniprot_accs: Collection[str], limit: int = 10_000, timeout: int = 1_800, batch_size: int = 10_000
-) -> dict[str, set[PdbResult]]:
+) -> PdbResults:
     """
     Search for PDB entries in UniProtKB accessions.
 
@@ -456,13 +627,20 @@ def search4pdb(
 
 
 def search4af(
-    uniprot_accs: Collection[str], limit: int = 10_000, timeout: int = 1_800, batch_size: int = 10_000
+    uniprot_accs: Collection[str],
+    min_sequence_length: int | None = None,
+    max_sequence_length: int | None = None,
+    limit: int = 10_000,
+    timeout: int = 1_800,
+    batch_size: int = 10_000,
 ) -> dict[str, set[str]]:
     """
     Search for AlphaFold entries in UniProtKB accessions.
 
     Args:
         uniprot_accs: UniProt accessions.
+        min_sequence_length: Minimum length of the canonical sequence.
+        max_sequence_length: Maximum length of the canonical sequence.
         limit: Maximum number of results to return.
         timeout: Timeout for the SPARQL query in seconds.
         batch_size: Size of batches to process the UniProt accessions.
@@ -474,7 +652,7 @@ def search4af(
     total = len(uniprot_accs)
     with tqdm(total=total, desc="Searching for AlphaFolds of uniprots", disable=total < batch_size, unit="acc") as pbar:
         for batch in batched(uniprot_accs, batch_size, strict=False):
-            sparql_query = _build_sparql_query_af(batch, limit)
+            sparql_query = _build_sparql_query_af(batch, min_sequence_length, max_sequence_length, limit)
             logger.info("Executing SPARQL query for AlphaFold: %s", sparql_query)
 
             raw_results = _execute_sparql_search(
@@ -639,12 +817,12 @@ def search4macromolecular_complexes(
 
 
 def search4interaction_partners(
-    uniprot_acc: str, excludes: set[str] | None = None, limit: int = 10_000, timeout: int = 1_800
+    uniprot_accession: str, excludes: set[str] | None = None, limit: int = 10_000, timeout: int = 1_800
 ) -> dict[str, set[str]]:
     """Search for interaction partners of a given UniProt accession using ComplexPortal database references.
 
     Args:
-        uniprot_acc: UniProt accession to search interaction partners for.
+        uniprot_accession: UniProt accession to search interaction partners for.
         excludes: Set of UniProt accessions to exclude from the results.
             For example already known interaction partners.
             If None then no complex members are excluded.
@@ -655,14 +833,143 @@ def search4interaction_partners(
         Dictionary with UniProt accessions of interaction partners as keys and sets of ComplexPortal entry IDs
         in which the interaction occurs as values.
     """
-    ucomplexes = search4macromolecular_complexes([uniprot_acc], limit=limit, timeout=timeout)
+    ucomplexes = search4macromolecular_complexes([uniprot_accession], limit=limit, timeout=timeout)
     hits: dict[str, set[str]] = {}
     if excludes is None:
         excludes = set()
     for ucomplex in ucomplexes:
         for member in ucomplex.members:
-            if member != uniprot_acc and member not in excludes:
+            if member != uniprot_accession and member not in excludes:
                 if member not in hits:
                     hits[member] = set()
                 hits[member].add(ucomplex.complex_id)
     return hits
+
+
+@dataclass(frozen=True)
+class UniprotDetails:
+    """Details of an UniProt entry.
+
+    Parameters:
+        uniprot_accession: UniProt accession.
+        uniprot_id: UniProt ID (mnemonic).
+        sequence_length: Length of the canonical sequence.
+        reviewed: Whether the entry is reviewed (Swiss-Prot) or unreviewed (TrEMBL).
+        protein_name: Recommended protein name.
+        taxon_id: NCBI Taxonomy ID of the organism.
+        taxon_name: Scientific name of the organism.
+    """
+
+    uniprot_accession: str
+    uniprot_id: str
+    sequence_length: int
+    reviewed: bool
+    protein_name: str
+    taxon_id: int
+    taxon_name: str
+
+
+def map_uniprot_accessions2uniprot_details(
+    uniprot_accessions: Collection[str], timeout: int = 1_800, batch_size: int = 1000
+) -> Generator[UniprotDetails]:
+    """Map UniProt accessions to UniProt details by querying the UniProt SPARQL endpoint.
+
+    Example:
+
+    SPARQL query to get details for 7 UniProt entries, run on [https://sparql.uniprot.org/sparql](https://sparql.uniprot.org/sparql).
+
+    ```sparql
+    PREFIX skos: <http://www.w3.org/2004/02/skos/core#>
+    PREFIX up:   <http://purl.uniprot.org/core/>
+    PREFIX rdf:  <http://www.w3.org/1999/02/22-rdf-syntax-ns#>
+
+    SELECT
+    (?ac AS ?uniprot_accession)
+    ?uniprot_id
+    (STRAFTER(STR(?organism), "taxonomy/") AS ?taxon_id)
+    ?taxon_name
+    ?reviewed
+    ?protein_name
+    (STRLEN(?sequence) AS ?seq_length)
+    WHERE {
+    # Input UniProt accessions
+    VALUES (?ac) { ("P05067") ("A6NGD5") ("O14627") ("P00697") ("P42284") ("A0A0B5AC95") ("A0A0S2Z4R0")}
+    BIND (IRI(CONCAT("http://purl.uniprot.org/uniprot/", ?ac)) AS ?protein)
+    ?protein a up:Protein .
+    ?protein up:mnemonic ?uniprot_id .
+    ?protein up:organism ?organism .
+    ?organism up:scientificName ?taxon_name .
+    ?protein up:reviewed ?reviewed .
+    ?protein up:recommendedName/up:fullName ?protein_name .
+    ?protein up:sequence ?isoform .
+    ?isoform a up:Simple_Sequence .
+    ?isoform rdf:value ?sequence .
+    BIND (IRI(STRBEFORE(REPLACE(
+        STR(?isoform), "http://purl.uniprot.org/isoforms/", "http://purl.uniprot.org/uniprot/"
+    ), "-")) AS ?ac_of_isoform)
+    FILTER(?ac_of_isoform = ?protein)
+    }
+    ```
+
+    Args:
+        uniprot_accessions: Iterable of UniProt accessions.
+        timeout: Timeout for the SPARQL query in seconds.
+        batch_size: Size of batches to process the UniProt accessions.
+
+    Yields:
+        UniprotDetails objects in random order.
+    """
+    select_clause = dedent("""\
+        (?ac AS ?uniprot_accession)
+        ?uniprot_id
+        (STRAFTER(STR(?organism), "taxonomy/") AS ?taxon_id)
+        ?taxon_name
+        ?reviewed
+        ?protein_name
+        (STRLEN(?sequence) AS ?seq_length)
+    """)
+    where_clause = dedent("""
+        ?protein up:mnemonic ?uniprot_id .
+        ?protein up:organism ?organism .
+        ?organism up:scientificName ?taxon_name .
+        ?protein up:reviewed ?reviewed .
+        OPTIONAL {
+        ?protein up:recommendedName/up:fullName ?protein_name .
+        }
+        ?protein up:sequence ?isoform .
+        ?isoform a up:Simple_Sequence .
+        ?isoform rdf:value ?sequence .
+        BIND (IRI(STRBEFORE(REPLACE(
+            STR(?isoform), "http://purl.uniprot.org/isoforms/", "http://purl.uniprot.org/uniprot/"
+        ), "-")) AS ?ac_of_isoform)
+        FILTER(?ac_of_isoform = ?protein)
+    """)
+    total = len(uniprot_accessions)
+    with tqdm(
+        total=total,
+        desc="Retrieving UniProt details",
+        disable=total < batch_size,
+        unit="acc",
+    ) as pbar:
+        for batch in batched(uniprot_accessions, batch_size, strict=False):
+            sparql_query = _build_sparql_generic_by_uniprot_accessions_query(
+                batch, select_clause, where_clause, limit=batch_size
+            )
+            logger.info("Executing SPARQL query for UniProt details: %s", sparql_query)
+            raw_results = _execute_sparql_search(
+                sparql_query=sparql_query,
+                timeout=timeout,
+            )
+            for raw_result in raw_results:
+                protein_name = raw_result.get("protein_name", {}).get("value", "")
+                result = UniprotDetails(
+                    uniprot_accession=raw_result["uniprot_accession"]["value"],
+                    uniprot_id=raw_result["uniprot_id"]["value"],
+                    sequence_length=int(raw_result["seq_length"]["value"]),
+                    reviewed=raw_result["reviewed"]["value"] == "true",
+                    protein_name=protein_name,
+                    taxon_id=int(raw_result["taxon_id"]["value"]),
+                    taxon_name=raw_result["taxon_name"]["value"],
+                )
+                yield result
+            pbar.update(len(batch))

protein_quest/utils.py

@@ -266,6 +266,7 @@ async def retrieve_files(
     cacher: Cacher | None = None,
     chunk_size: int = 524288,  # 512 KiB
     gzip_files: bool = False,
+    raise_for_not_found: bool = True,
 ) -> list[Path]:
     """Retrieve files from a list of URLs and save them to a directory.
 
@@ -279,6 +280,9 @@ async def retrieve_files(
         cacher: An optional cacher to use for caching files.
         chunk_size: The size of each chunk to read from the response.
         gzip_files: Whether to gzip the downloaded files.
+            This requires the server can send gzip encoded content.
+        raise_for_not_found: Whether to raise an error for HTTP 404 errors.
+            If false then function does not returns Path for which url gave HTTP 404 error and logs as debug message.
 
     Returns:
         A list of paths to the downloaded files.
@@ -295,11 +299,12 @@ async def retrieve_files(
                 cacher=cacher,
                 chunk_size=chunk_size,
                 gzip_files=gzip_files,
+                raise_for_not_found=raise_for_not_found,
             )
             for url, filename in urls
         ]
-        files: list[Path] = await tqdm.gather(*tasks, desc=desc)
-        return files
+        raw_files: list[Path | None] = await tqdm.gather(*tasks, desc=desc)
+        return [f for f in raw_files if f is not None]
 
 
 class InvalidContentEncodingError(aiohttp.ClientResponseError):
@@ -314,7 +319,8 @@ async def _retrieve_file(
     cacher: Cacher | None = None,
     chunk_size: int = 524288,  # 512 KiB
     gzip_files: bool = False,
-) -> Path:
+    raise_for_not_found=True,
+) -> Path | None:
     """Retrieve a single file from a URL and save it to a specified path.
 
     Args:
@@ -325,6 +331,9 @@ async def _retrieve_file(
         cacher: An optional cacher to use for caching files.
         chunk_size: The size of each chunk to read from the response.
         gzip_files: Whether to gzip the downloaded file.
+            This requires the server can send gzip encoded content.
+        raise_for_not_found: Whether to raise an error for HTTP 404 errors.
+            If false then function returns None on HTTP 404 errors and logs as debug message.
 
     Returns:
         The path to the saved file.
@@ -348,6 +357,9 @@ async def _retrieve_file(
         semaphore,
         session.get(url, headers=headers, auto_decompress=auto_decompress) as resp,
     ):
+        if not raise_for_not_found and resp.status == 404:
+            logger.debug(f"File not found at {url}, skipping download.")
+            return None
         resp.raise_for_status()
         if gzip_files and resp.headers.get("Content-Encoding") != "gzip":
             msg = f"Server did not send gzip encoded content for {url}, can not save as gzipped file."