protein-quest 0.5.1__py3-none-any.whl → 0.7.0__py3-none-any.whl

protein_quest/{pdbe/io.py → structure.py}

@@ -1,51 +1,29 @@
-"""Module for structure file input/output."""
+"""Module for querying and modifying [gemmi structures][gemmi.Structure]."""
 
-import gzip
 import logging
-from collections.abc import Generator, Iterable
+from collections.abc import Iterable
 from datetime import UTC, datetime
 from pathlib import Path
 
 import gemmi
 
 from protein_quest.__version__ import __version__
+from protein_quest.io import read_structure, split_name_and_extension, write_structure
 from protein_quest.utils import CopyMethod, copyfile
 
 logger = logging.getLogger(__name__)
 
-# TODO remove once v0.7.4 of gemmi is released,
-# as uv pip install git+https://github.com/project-gemmi/gemmi.git installs 0.7.4.dev0 which does not print leaks
-# Swallow gemmi leaked function warnings
-gemmi.set_leak_warnings(False)
 
-
-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.
+def find_chain_in_model(model: gemmi.Model, wanted_chain: str) -> gemmi.Chain | None:
+    """Find a chain in a model.
 
     Args:
-        file: Path to the input mmCIF/pdb file.
-        chain: Chain to count residues of.
+        model: The gemmi model to search in.
+        wanted_chain: The chain identifier to search for.
 
     Returns:
-        The number of residues in the specified chain.
+        The found chain or None if not found.
     """
-    structure = gemmi.read_structure(str(file))
-    gchain = find_chain_in_structure(structure, 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_structure(structure: gemmi.Structure, wanted_chain: str) -> gemmi.Chain | None:
-    for model in structure:
-        chain = find_chain_in_model(model, wanted_chain)
-        if chain is not None:
-            return chain
-    return None
-
-
-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,
@@ -57,106 +35,39 @@ def find_chain_in_model(model: gemmi.Model, wanted_chain: str) -> gemmi.Chain |
     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"):
-        # do not write chem_comp so it is viewable by molstar
-        # see https://github.com/project-gemmi/gemmi/discussions/362
-        doc = structure.make_mmcif_document(gemmi.MmcifOutputGroups(True, chem_comp=False))
-        doc.write_file(str(path))
-    elif path.name.endswith(".cif.gz"):
-        doc = structure.make_mmcif_document(gemmi.MmcifOutputGroups(True, chem_comp=False))
-        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.
+def find_chain_in_structure(structure: gemmi.Structure, wanted_chain: str) -> gemmi.Chain | None:
+    """Find a chain in a structure.
 
     Args:
-        root: The root directory to search in.
-        pdb_id: The PDB ID to locate.
+        structure: The gemmi structure to search in.
+        wanted_chain: The chain identifier to search for.
 
     Returns:
-        The path to the located structure file.
-
-    Raises:
-        FileNotFoundError: If no structure file is found for the given PDB ID.
+        The found chain or None if not found.
     """
-    exts = [".cif.gz", ".cif", ".pdb.gz", ".pdb", ".ent", ".ent.gz"]
-    for ext in exts:
-        candidates = (
-            root / f"{pdb_id}{ext}",
-            root / f"{pdb_id.lower()}{ext}",
-            root / f"{pdb_id.upper()}{ext}",
-            root / f"pdb{pdb_id.lower()}{ext}",
-        )
-        for candidate in candidates:
-            if candidate.exists():
-                return candidate
-    msg = f"No structure file found for {pdb_id} in {root}"
-    raise FileNotFoundError(msg)
+    for model in structure:
+        chain = find_chain_in_model(model, wanted_chain)
+        if chain is not None:
+            return chain
+    return None
 
 
-def glob_structure_files(input_dir: Path) -> Generator[Path]:
-    """Glob for structure files in a directory.
+def nr_residues_in_chain(file: Path, chain: str = "A") -> int:
+    """Returns the number of residues in a specific chain from a structure file.
 
     Args:
-        input_dir: The input directory to search for structure files.
+        file: Path to the input structure file.
+        chain: Chain to count residues of.
 
-    Yields:
-        Paths to the found structure files.
+    Returns:
+        The number of residues in the specified chain.
     """
-    for ext in [".cif.gz", ".cif", ".pdb.gz", ".pdb"]:
-        yield from input_dir.glob(f"*{ext}")
-
-
-class ChainNotFoundError(IndexError):
-    """Exception raised when a chain is not found in a structure."""
-
-    def __init__(self, chain: str, file: Path | str, available_chains: Iterable[str]):
-        super().__init__(f"Chain {chain} not found in {file}. Available chains are: {available_chains}")
-        self.chain_id = chain
-        self.file = file
+    structure = read_structure(file)
+    gchain = find_chain_in_structure(structure, chain)
+    if gchain is None:
+        logger.warning("Chain %s not found in %s. Returning 0.", chain, file)
+        return 0
+    return len(gchain)
 
 
 def _dedup_helices(structure: gemmi.Structure):
@@ -198,18 +109,34 @@ def _add_provenance_info(structure: gemmi.Structure, chain2keep: str, out_chain:
 
 
 def chains_in_structure(structure: gemmi.Structure) -> set[gemmi.Chain]:
-    """Get a list of chains in a structure."""
+    """Get a list of chains in a structure.
+
+    Args:
+        structure: The gemmi structure to get chains from.
+
+    Returns:
+        A set of chains in the structure.
+    """
     return {c for model in structure for c in model}
 
 
-def write_single_chain_pdb_file(
+class ChainNotFoundError(IndexError):
+    """Exception raised when a chain is not found in a structure."""
+
+    def __init__(self, chain: str, file: Path | str, available_chains: Iterable[str]):
+        super().__init__(f"Chain {chain} not found in {file}. Available chains are: {available_chains}")
+        self.chain_id = chain
+        self.file = file
+
+
+def write_single_chain_structure_file(
     input_file: Path,
     chain2keep: str,
     output_dir: Path,
     out_chain: str = "A",
     copy_method: CopyMethod = "copy",
 ) -> Path:
-    """Write a single chain from a mmCIF/pdb file to a new mmCIF/pdb file.
+    """Write a single chain from a structure file to a new structure file.
 
     Also
 
@@ -226,14 +153,14 @@ def write_single_chain_pdb_file(
     ```
 
     Args:
-        input_file: Path to the input mmCIF/pdb file.
+        input_file: Path to the input structure file.
         chain2keep: The chain to keep.
         output_dir: Directory to save the output file.
         out_chain: The chain identifier for the output file.
         copy_method: How to copy when no changes are needed to output file.
 
     Returns:
-        Path to the output mmCIF/pdb file
+        Path to the output structure file
 
     Raises:
         FileNotFoundError: If the input file does not exist.
@@ -241,7 +168,7 @@ def write_single_chain_pdb_file(
     """
 
     logger.debug(f"chain2keep: {chain2keep}, out_chain: {out_chain}")
-    structure = gemmi.read_structure(str(input_file))
+    structure = read_structure(input_file)
     structure.setup_entities()
 
     chain = find_chain_in_structure(structure, chain2keep)
@@ -249,7 +176,7 @@ def write_single_chain_pdb_file(
     if chain is None:
         raise ChainNotFoundError(chain2keep, input_file, chainnames_in_structure)
     chain_name = chain.name
-    name, extension = _split_name_and_extension(input_file.name)
+    name, extension = split_name_and_extension(input_file.name)
     output_file = output_dir / f"{name}_{chain_name}2{out_chain}{extension}"
 
     if output_file.exists():
@@ -279,3 +206,27 @@ def write_single_chain_pdb_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