protein-quest 0.3.0__py3-none-any.whl

protein_quest/pdbe/io.py

@@ -0,0 +1,185 @@
+"""Module for structure file input/output."""
+
+import gzip
+import logging
+from collections.abc import Generator
+from pathlib import Path
+
+import gemmi
+
+from protein_quest import __version__
+
+logger = logging.getLogger(__name__)
+
+
+def nr_residues_in_chain(file: Path | str, chain: str = "A") -> int:
+    """Returns the number of residues in a specific chain from a mmCIF/pdb file.
+
+    Args:
+        file: Path to the input mmCIF/pdb file.
+        chain: Chain to count residues of.
+
+    Returns:
+        The number of residues in the specified chain.
+    """
+    structure = gemmi.read_structure(str(file))
+    model = structure[0]
+    gchain = find_chain_in_model(model, chain)
+    if gchain is None:
+        logger.warning("Chain %s not found in %s. Returning 0.", chain, file)
+        return 0
+    return len(gchain)
+
+
+def find_chain_in_model(model: gemmi.Model, wanted_chain: str) -> gemmi.Chain | None:
+    chain = model.find_chain(wanted_chain)
+    if chain is None:
+        # For chain A in 4v92 the find_chain method returns None,
+        # however it is prefixed with 'B',
+        # so we try again as last char of chain name
+        mchains = [c for c in model if c.name.endswith(wanted_chain)]
+        if mchains:
+            return mchains[0]
+    return chain
+
+
+def write_structure(structure: gemmi.Structure, path: Path):
+    """Write a gemmi structure to a file.
+
+    Args:
+        structure: The gemmi structure to write.
+        path: The file path to write the structure to.
+            The format depends on the file extension.
+            Supported extensions are .pdb, .pdb.gz, .cif, .cif.gz.
+
+    Raises:
+        ValueError: If the file extension is not supported.
+    """
+    if path.name.endswith(".pdb"):
+        body: str = structure.make_pdb_string()
+        path.write_text(body)
+    elif path.name.endswith(".pdb.gz"):
+        body: str = structure.make_pdb_string()
+        with gzip.open(path, "wt") as f:
+            f.write(body)
+    elif path.name.endswith(".cif"):
+        doc = structure.make_mmcif_document()
+        doc.write_file(str(path))
+    elif path.name.endswith(".cif.gz"):
+        doc = structure.make_mmcif_document()
+        cif_str = doc.as_string()
+        with gzip.open(path, "wt") as f:
+            f.write(cif_str)
+    else:
+        msg = f"Unsupported file extension in {path.name}. Supported extensions are .pdb, .pdb.gz, .cif, .cif.gz"
+        raise ValueError(msg)
+
+
+def _split_name_and_extension(name: str) -> tuple[str, str]:
+    # 1234.pdb -> (1234, .pdb)
+    # 1234.pdb.gz -> (1234, .pdb.gz)
+    # 1234.cif -> (1234, .cif)
+    # 1234.cif.gz -> (1234, .cif.gz)
+    if name.endswith(".pdb.gz"):
+        return name.replace(".pdb.gz", ""), ".pdb.gz"
+    if name.endswith(".cif.gz"):
+        return name.replace(".cif.gz", ""), ".cif.gz"
+    if name.endswith(".pdb"):
+        return name.replace(".pdb", ""), ".pdb"
+    if name.endswith(".cif"):
+        return name.replace(".cif", ""), ".cif"
+
+    msg = f"Unknown file extension in {name}. Supported extensions are .pdb, .pdb.gz, .cif, .cif.gz"
+    raise ValueError(msg)
+
+
+def locate_structure_file(root: Path, pdb_id: str) -> Path:
+    """Locate a structure file for a given PDB ID in the specified directory.
+
+    Args:
+        root: The root directory to search in.
+        pdb_id: The PDB ID to locate.
+
+    Returns:
+        The path to the located structure file.
+
+    Raises:
+        FileNotFoundError: If no structure file is found for the given PDB ID.
+    """
+    exts = [".cif.gz", ".cif", ".pdb.gz", ".pdb"]
+    # files downloaded from https://www.ebi.ac.uk/pdbe/ website
+    # have file names like pdb6t5y.ent or pdb6t5y.ent.gz for a PDB formatted file.
+    # TODO support pdb6t5y.ent or pdb6t5y.ent.gz file names
+    for ext in exts:
+        candidate = root / f"{pdb_id.lower()}{ext}"
+        if candidate.exists():
+            return candidate
+    msg = f"No structure file found for {pdb_id} in {root}"
+    raise FileNotFoundError(msg)
+
+
+def glob_structure_files(input_dir: Path) -> Generator[Path]:
+    """Glob for structure files in a directory.
+
+    Args:
+        input_dir: The input directory to search for structure files.
+
+    Yields:
+        Paths to the found structure files.
+    """
+    for ext in [".cif.gz", ".cif", ".pdb.gz", ".pdb"]:
+        yield from input_dir.glob(f"*{ext}")
+
+
+def write_single_chain_pdb_file(
+    input_file: Path, chain2keep: str, output_dir: Path, out_chain: str = "A"
+) -> Path | None:
+    """Write a single chain from a mmCIF/pdb file to a new mmCIF/pdb file.
+
+    Args:
+        input_file: Path to the input mmCIF/pdb file.
+        chain2keep: The chain to keep.
+        output_dir: Directory to save the output file.
+        out_chain: The chain identifier for the output file.
+
+    Returns:
+        Path to the output mmCIF/pdb file or None if not created.
+    """
+
+    structure = gemmi.read_structure(str(input_file))
+    model = structure[0]
+
+    # Only count residues of polymer
+    model.remove_ligands_and_waters()
+
+    chain = find_chain_in_model(model, chain2keep)
+    if chain is None:
+        logger.warning(
+            "Chain %s not found in %s. Skipping.",
+            chain2keep,
+            input_file,
+        )
+        return None
+    name, extension = _split_name_and_extension(input_file.name)
+    output_file = output_dir / f"{name}_{chain.name}2{out_chain}{extension}"
+
+    new_structure = gemmi.Structure()
+    new_structure.resolution = structure.resolution
+    new_id = structure.name + f"{chain2keep}2{out_chain}"
+    new_structure.name = new_id
+    new_structure.info["_entry.id"] = new_id
+    new_title = f"From {structure.info['_entry.id']} chain {chain2keep} to {out_chain}"
+    new_structure.info["_struct.title"] = new_title
+    new_structure.info["_struct_keywords.pdbx_keywords"] = new_title.upper()
+    new_si = gemmi.SoftwareItem()
+    new_si.classification = gemmi.SoftwareItem.Classification.DataExtraction
+    new_si.name = "protein-quest"
+    new_si.version = str(__version__)
+    new_structure.meta.software.append(new_si)
+    new_model = gemmi.Model(1)
+    chain.name = out_chain
+    new_model.add_chain(chain)
+    new_structure.add_model(new_model)
+    write_structure(new_structure, output_file)
+
+    return output_file

protein_quest/taxonomy.py

@@ -0,0 +1,139 @@
+"""Module for searching taxon information from UniProt."""
+
+import csv
+import gzip
+import logging
+from dataclasses import dataclass
+from typing import Literal, get_args
+
+from aiohttp.client import ClientResponse
+from aiohttp_retry import RetryClient
+from cattrs.gen import make_dict_structure_fn, override
+from cattrs.preconf.orjson import make_converter
+from yarl import URL
+
+from protein_quest.go import TextIOWrapper
+from protein_quest.utils import friendly_session
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True, slots=True)
+class Taxon:
+    taxon_id: str
+    scientific_name: str
+    rank: str
+    common_name: str | None = None
+    other_names: set[str] | None = None
+
+
+@dataclass(frozen=True, slots=True)
+class SearchTaxonResponse:
+    results: list[Taxon]
+
+
+converter = make_converter()
+
+converter.register_structure_hook(
+    Taxon,
+    make_dict_structure_fn(
+        Taxon,
+        converter,
+        taxon_id=override(rename="taxonId"),
+        scientific_name=override(rename="scientificName"),
+        common_name=override(rename="commonName"),
+        other_names=override(rename="otherNames"),
+    ),
+)
+
+SearchField = Literal["tax_id", "scientific", "common", "parent"]
+search_fields: set[SearchField | None] = set(get_args(SearchField)) | {None}
+
+
+def _get_next_page(response: ClientResponse) -> URL | str | None:
+    next_page = response.links.getone("next", None)
+    if next_page is None:
+        return None
+    return next_page.getone("url", None)
+
+
+async def _fetch_page(url: URL | str, session: RetryClient) -> tuple[list[Taxon], URL | str | None]:
+    async with session.get(url) as response:
+        response.raise_for_status()
+        gzipped_raw_data = await response.read()
+    next_page = _get_next_page(response)
+    raw_data = gzip.decompress(gzipped_raw_data)
+    taxons = converter.loads(raw_data, SearchTaxonResponse).results
+    return taxons, next_page
+
+
+async def search_taxon(query: str, field: SearchField | None = None, limit: int = 100) -> list[Taxon]:
+    """Search for taxon information in UniProt.
+
+    Uses <https://www.uniprot.org/taxonomy?query=*>.
+
+    Args:
+        query: Search query for the taxon.
+        field: Field to search in.
+            If None, searches in all fields.
+            If "tax_id" then searches by taxon ID.
+            If "parent" then given a parent taxon ID returns all its children.
+            For example, if the parent taxon ID is 9606 (Human), it will return Neanderthal and Denisovan.
+        limit: Maximum number of results to return.
+
+    Returns:
+        List of Taxon objects matching the search query.
+
+    Raises:
+        ValueError: If the search field is invalid.
+    """
+    # https://rest.uniprot.org/taxonomy/search?compressed=true&format=json&query=%28Mouse%29&size=500
+    page_limit = 100
+    if field not in search_fields:
+        msg = f"Invalid search field: {field}. Must be one of {search_fields}."
+        raise ValueError(msg)
+    if field is not None:
+        # ((common:"house+mouse"))
+        query = f'(({field}:"{query}"))'
+    params = {"query": query, "limit": str(page_limit), "compressed": "true", "format": "json"}
+    url = URL("https://rest.uniprot.org/taxonomy/search").with_query(params)
+    logger.debug("Fetching uniprot taxonomy from %s with params %s", url, params)
+    async with friendly_session() as session:
+        # Fetch first page
+        taxons, next_page = await _fetch_page(url, session)
+        if len(taxons) >= limit:
+            return taxons[:limit]
+        if next_page is None:
+            return taxons
+
+        # Fetch next pages
+        while next_page:
+            logger.debug("Fetching next page of uniprot taxonomy from %s", next_page)
+            next_data, next_page = await _fetch_page(next_page, session)
+            taxons.extend(next_data)
+            if len(taxons) >= limit:
+                return taxons[:limit]
+            if next_page is None:
+                return taxons
+    return taxons
+
+
+def _write_taxonomy_csv(taxons: list[Taxon], output_csv: TextIOWrapper) -> None:
+    """Write taxon information to a CSV file.
+
+    Args:
+        taxons: List of Taxon objects to write to the CSV file.
+        output_csv: File object for the output CSV file.
+    """
+    writer = csv.writer(output_csv)
+    writer.writerow(["taxon_id", "scientific_name", "common_name", "rank", "other_names"])
+    for taxon in taxons:
+        writer.writerow(
+            [
+                taxon.taxon_id,
+                taxon.scientific_name,
+                taxon.common_name,
+                taxon.rank,
+                ";".join(taxon.other_names) if taxon.other_names else "",
+            ]
+        )