protein-quest 0.9.0__py3-none-any.whl

protein_quest/io.py

@@ -0,0 +1,350 @@
+"""Module for structure file input/output."""
+
+import gzip
+import logging
+import shutil
+import tempfile
+from collections.abc import Generator, Iterable
+from io import StringIO
+from pathlib import Path
+from typing import Literal, get_args
+from urllib.request import urlopen
+
+import gemmi
+from mmcif.api.DictionaryApi import DictionaryApi
+from mmcif.io.BinaryCifReader import BinaryCifReader
+from mmcif.io.BinaryCifWriter import BinaryCifWriter
+from mmcif.io.PdbxReader import PdbxReader
+from mmcif.io.PdbxWriter import PdbxWriter
+
+from protein_quest.utils import CopyMethod, copyfile, user_cache_root_dir
+
+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)
+
+
+StructureFileExtensions = Literal[".pdb", ".pdb.gz", ".ent", ".ent.gz", ".cif", ".cif.gz", ".bcif", ".bcif.gz"]
+"""Type of supported structure file extensions."""
+valid_structure_file_extensions: set[str] = set(get_args(StructureFileExtensions))
+"""Set of valid structure file extensions."""
+
+
+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.
+            See [StructureFileExtensions][protein_quest.io.StructureFileExtensions]
+            for supported extensions.
+
+    Raises:
+        ValueError: If the file extension is not supported.
+    """
+    if path.name.endswith(".pdb") or path.name.endswith(".ent"):
+        body: str = structure.make_pdb_string()
+        path.write_text(body)
+    elif path.name.endswith(".pdb.gz") or path.name.endswith(".ent.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)
+    elif path.name.endswith(".bcif"):
+        structure2bcif(structure, path)
+    elif path.name.endswith(".bcif.gz"):
+        structure2bcifgz(structure, path)
+    else:
+        msg = f"Unsupported file extension in {path.name}. Supported extensions are: {valid_structure_file_extensions}"
+        raise ValueError(msg)
+
+
+def read_structure(file: Path) -> gemmi.Structure:
+    """Read a structure from a file.
+
+    Args:
+        file: Path to the input structure file.
+            See [StructureFileExtensions][protein_quest.io.StructureFileExtensions]
+            for supported extensions.
+
+    Returns:
+        A gemmi Structure object representing the structure in the file.
+    """
+    if file.name.endswith(".bcif"):
+        return bcif2structure(file)
+    if file.name.endswith(".bcif.gz"):
+        return bcifgz2structure(file)
+    return gemmi.read_structure(str(file))
+
+
+def bcif2cif(bcif_file: Path) -> str:
+    """Convert a binary CIF (bcif) file to a CIF string.
+
+    Args:
+        bcif_file: Path to the binary CIF file.
+
+    Returns:
+        A string containing the CIF representation of the structure.
+    """
+    reader = BinaryCifReader()
+    container = reader.deserialize(str(bcif_file))
+    capture = StringIO()
+    writer = PdbxWriter(capture)
+    writer.write(container)
+    return capture.getvalue()
+
+
+def bcifgz2structure(bcif_gz_file: Path) -> gemmi.Structure:
+    """Read a binary CIF (bcif) gzipped file and return a gemmi Structure object.
+
+    This is slower than other formats because gemmi does not support reading bcif files directly.
+    So we first gunzip the file to a temporary location, convert it to a cif string using mmcif package,
+    and then read the cif string using gemmi.
+
+    Args:
+        bcif_gz_file: Path to the binary CIF gzipped file.
+
+    Returns:
+        A gemmi Structure object representing the structure in the bcif.gz file.
+    """
+    with tempfile.NamedTemporaryFile(suffix=".bcif", delete=True) as tmp_bcif:
+        tmp_path = Path(tmp_bcif.name)
+        gunzip_file(bcif_gz_file, output_file=tmp_path, keep_original=True)
+        return bcif2structure(tmp_path)
+
+
+def bcif2structure(bcif_file: Path) -> gemmi.Structure:
+    """Read a binary CIF (bcif) file and return a gemmi Structure object.
+
+    This is slower than other formats because gemmi does not support reading bcif files directly.
+    So we convert it to a cif string first using mmcif package and then read the cif string using gemmi.
+
+    Args:
+        bcif_file: Path to the binary CIF file.
+
+    Returns:
+        A gemmi Structure object representing the structure in the bcif file.
+    """
+    cif_content = bcif2cif(bcif_file)
+    doc = gemmi.cif.read_string(cif_content)
+    block = doc.sole_block()
+    return gemmi.make_structure_from_block(block)
+
+
+def _initialize_dictionary_api(containers) -> DictionaryApi:
+    dict_local = user_cache_root_dir() / "mmcif_pdbx_v5_next.dic"
+    if not dict_local.exists():
+        dict_url = "https://raw.githubusercontent.com/wwpdb-dictionaries/mmcif_pdbx/master/dist/mmcif_pdbx_v5_next.dic"
+        logger.info("Downloading mmcif dictionary from %s to %s", dict_url, dict_local)
+        dict_local.parent.mkdir(parents=True, exist_ok=True)
+        with dict_local.open("wb") as f, urlopen(dict_url) as response:  # noqa: S310 url is hardcoded and https
+            f.write(response.read())
+    return DictionaryApi(containerList=containers, consolidate=True)
+
+
+def structure2bcif(structure: gemmi.Structure, bcif_file: Path):
+    """Write a gemmi Structure object to a binary CIF (bcif) file.
+
+    This is slower than other formats because gemmi does not support writing bcif files directly.
+    So we convert it to a cif string first using gemmi and then convert cif to bcif using mmcif package.
+
+    Args:
+        structure: The gemmi Structure object to write.
+        bcif_file: Path to the output binary CIF file.
+    """
+    doc = structure.make_mmcif_document(gemmi.MmcifOutputGroups(True, chem_comp=False))
+    containers = []
+    with StringIO(doc.as_string()) as sio:
+        reader = PdbxReader(sio)
+        reader.read(containers)
+    dict_api = _initialize_dictionary_api(containers)
+    writer = BinaryCifWriter(dictionaryApi=dict_api)
+    writer.serialize(str(bcif_file), containers)
+
+
+def gunzip_file(gz_file: Path, output_file: Path | None = None, keep_original: bool = True) -> Path:
+    """Unzip a .gz file.
+
+    Args:
+        gz_file: Path to the .gz file.
+        output_file: Optional path to the output unzipped file. If None, the .gz suffix is removed from gz_file.
+        keep_original: Whether to keep the original .gz file. Default is True.
+
+    Returns:
+        Path to the unzipped file.
+
+    Raises:
+        ValueError: If output_file is None and gz_file does not end with .gz.
+    """
+    if output_file is None and not gz_file.name.endswith(".gz"):
+        msg = f"If output_file is not provided, {gz_file} must end with .gz"
+        raise ValueError(msg)
+    out_file = output_file or gz_file.with_suffix("")
+    with gzip.open(gz_file, "rb") as f_in, out_file.open("wb") as f_out:
+        shutil.copyfileobj(f_in, f_out)
+    if not keep_original:
+        gz_file.unlink()
+    return out_file
+
+
+def structure2bcifgz(structure: gemmi.Structure, bcif_gz_file: Path):
+    """Write a gemmi Structure object to a binary CIF gzipped (bcif.gz) file.
+
+    This is slower than other formats because gemmi does not support writing bcif files directly.
+    So we convert it to a cif string first using gemmi and then convert cif to bcif using mmcif package.
+    Finally, we gzip the bcif file.
+
+    Args:
+        structure: The gemmi Structure object to write.
+        bcif_gz_file: Path to the output binary CIF gzipped file.
+    """
+    with tempfile.NamedTemporaryFile(suffix=".bcif", delete=True) as tmp_bcif:
+        tmp_path = Path(tmp_bcif.name)
+        structure2bcif(structure, tmp_path)
+        with tmp_path.open("rb") as f_in, gzip.open(bcif_gz_file, "wb") as f_out:
+            shutil.copyfileobj(f_in, f_out)
+
+
+def convert_to_cif_files(
+    input_files: Iterable[Path], output_dir: Path, copy_method: CopyMethod
+) -> Generator[tuple[Path, Path]]:
+    """Convert structure files to .cif format.
+
+    Args:
+        input_files: Iterable of structure files to convert.
+        output_dir: Directory to save the converted .cif files.
+        copy_method: How to copy when no changes are needed to output file.
+
+    Yields:
+        A tuple of the input file and the output file.
+    """
+    for input_file in input_files:
+        output_file = convert_to_cif_file(input_file, output_dir, copy_method)
+        yield input_file, output_file
+
+
+def convert_to_cif_file(input_file: Path, output_dir: Path, copy_method: CopyMethod) -> Path:
+    """Convert a single structure file to .cif format.
+
+    Args:
+        input_file: The structure file to convert.
+            See [StructureFileExtensions][protein_quest.io.StructureFileExtensions]
+            for supported extensions.
+        output_dir: Directory to save the converted .cif file.
+        copy_method: How to copy when no changes are needed to output file.
+
+    Returns:
+        Path to the converted .cif file.
+    """
+    name, extension = split_name_and_extension(input_file.name)
+    output_file = output_dir / f"{name}.cif"
+    if output_file.exists():
+        logger.info("Output file %s already exists for input file %s. Skipping.", output_file, input_file)
+    elif extension in {".pdb", ".pdb.gz", ".ent", ".ent.gz"}:
+        structure = read_structure(input_file)
+        write_structure(structure, output_file)
+    elif extension == ".cif":
+        logger.info("File %s is already in .cif format, copying to %s", input_file, output_dir)
+        copyfile(input_file, output_file, copy_method)
+    elif extension == ".cif.gz":
+        gunzip_file(input_file, output_file=output_file, keep_original=True)
+    elif extension == ".bcif":
+        with output_file.open("w") as f:
+            f.write(bcif2cif(input_file))
+    else:
+        msg = (
+            f"Unsupported file extension {extension} in {input_file}. "
+            f"Supported extensions are {valid_structure_file_extensions}."
+        )
+        raise ValueError(msg)
+    return output_file
+
+
+def split_name_and_extension(name: str) -> tuple[str, str]:
+    """Split a filename into its name and extension.
+
+    `.gz` is considered part of the extension if present.
+
+    Examples:
+        Some example usages.
+
+        >>> from protein_quest.pdbe.io import split_name_and_extension
+        >>> split_name_and_extension("1234.pdb")
+        ('1234', '.pdb')
+        >>> split_name_and_extension("1234.pdb.gz")
+        ('1234', '.pdb.gz')
+
+    Args:
+        name: The filename to split.
+
+    Returns:
+        A tuple containing the name and the extension.
+    """
+    ext = ""
+    if name.endswith(".gz"):
+        ext = ".gz"
+        name = name.removesuffix(".gz")
+    i = name.rfind(".")
+    if 0 < i < len(name) - 1:
+        ext = name[i:] + ext
+        name = name[:i]
+    return name, ext
+
+
+def locate_structure_file(root: Path, pdb_id: str) -> Path:
+    """Locate a structure file for a given PDB ID in the specified directory.
+
+    Uses [StructureFileExtensions][protein_quest.io.StructureFileExtensions] as potential extensions.
+    Also tries different casing of the PDB ID.
+
+    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.
+    """
+    for ext in valid_structure_file_extensions:
+        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)
+
+
+def glob_structure_files(input_dir: Path) -> Generator[Path]:
+    """Glob for structure files in a directory.
+
+    Uses [StructureFileExtensions][protein_quest.io.StructureFileExtensions] as valid extensions.
+    Does not search recursively.
+
+    Args:
+        input_dir: The input directory to search for structure files.
+
+    Yields:
+        Paths to the found structure files.
+    """
+    for ext in valid_structure_file_extensions:
+        yield from input_dir.glob(f"*{ext}")

protein_quest/mcp_server.py

@@ -0,0 +1,256 @@
+"""MCP server for protein-quest.
+
+Can be run with:
+
+```shell
+# for development
+fastmcp dev src/protein_quest/mcp_server.py
+# or from inspector
+npx @modelcontextprotocol/inspector
+# tranport type: stdio
+# comand: protein-quest
+# arguments: mcp
+
+# or with server and inspector
+protein-quest mcp --transport streamable-http
+# in another shell
+npx @modelcontextprotocol/inspector
+# transport type: streamable http
+# URL: http://127.0.0.1:8000/mcp
+
+# or with copilot in VS code
+# ctrl + shift + p
+# mcp: add server...
+# Choose STDIO
+# command: uv run protein-quest mcp
+# id: protein-quest
+```
+
+Examples:
+
+   - What are the PDBe structures for `A8MT69` uniprot accession?
+
+"""
+
+from collections.abc import Mapping
+from pathlib import Path
+from textwrap import dedent
+from typing import Annotated
+
+from fastmcp import FastMCP
+from pydantic import Field
+
+from protein_quest.alphafold.confidence import ConfidenceFilterQuery, ConfidenceFilterResult, filter_file_on_confidence
+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, 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,
+    structure2uniprot_accessions,
+    write_single_chain_structure_file,
+)
+from protein_quest.taxonomy import search_taxon
+from protein_quest.uniprot import (
+    PdbResult,
+    Query,
+    search4af,
+    search4emdb,
+    search4macromolecular_complexes,
+    search4pdb,
+    search4uniprot,
+)
+
+mcp = FastMCP("protein-quest")
+
+# do not want to make dataclasses in non-mcp code into Pydantic models,
+# so we use Annotated here to add description on roots.
+
+
+@mcp.tool
+def search_uniprot(
+    uniprot_query: Annotated[Query, Field(description=Query.__doc__)],
+    limit: Annotated[int, Field(gt=0, description="Limit the number of uniprot accessions returned")] = 100,
+) -> set[str]:
+    """Search UniProt for proteins matching the given query."""
+    return search4uniprot(uniprot_query, limit=limit)
+
+
+@mcp.tool
+def search_pdb(
+    uniprot_accs: set[str],
+    limit: Annotated[int, Field(gt=0, description="Limit the number of entries returned")] = 100,
+) -> Annotated[
+    dict[str, set[PdbResult]],
+    Field(
+        description=dedent(f"""\
+            Dictionary with protein IDs as keys and sets of PDB results as values.
+            A PDB result is {PdbResult.__doc__}""")
+    ),
+]:
+    """Search PDBe structures for given uniprot accessions."""
+    return search4pdb(uniprot_accs, limit=limit)
+
+
+@mcp.tool
+async def fetch_pdbe_structures(pdb_ids: set[str], save_dir: Path) -> Mapping[str, Path]:
+    """Fetch the PDBe structures for given PDB IDs.
+
+    Args:
+        pdb_ids: A set of PDB IDs.
+        save_dir: The directory to save the fetched files.
+
+    Returns:
+        A mapping of PDB ID to the path of the fetched structure file.
+    """
+    return await pdbe_fetch(pdb_ids, save_dir)
+
+
+@mcp.tool
+def extract_single_chain_from_structure(
+    input_file: Path,
+    chain2keep: str,
+    output_dir: Path,
+    out_chain: str = "A",
+) -> Path:
+    """
+    Extract a single chain from a structure (mmCIF or pdb) file and write to a new file.
+
+    Args:
+        input_file: Path to the input structure (mmCIF or 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 structure (mmCIF or pdb) file
+    """
+    return write_single_chain_structure_file(input_file, chain2keep, output_dir, out_chain)
+
+
+@mcp.tool
+def list_structure_files(path: Path) -> list[Path]:
+    """List structure files (.pdb, .pdb.gz, .cif, .cif.gz, .bcif) in the specified directory."""
+    return list(glob_structure_files(path))
+
+
+# TODO replace remaining decorators with wrapper if tool does single function call
+# so we do not have to replicate docstring,
+# minor con is that it does not show up in api docs
+mcp.tool(nr_residues_in_chain)
+mcp.tool(search_taxon)
+mcp.tool(search_gene_ontology_term)
+
+
+@mcp.tool
+def search_alphafolds(
+    uniprot_accs: set[str],
+    limit: Annotated[int, Field(gt=0, description="Limit the number of entries returned")] = 100,
+) -> Annotated[
+    set[str],
+    Field(description="Set of uniprot accessions which have an AlphaFold entry"),
+]:
+    """Search for AlphaFold entries in UniProtKB accessions."""
+    # 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)
+    return {k for k, v in results.items() if v}
+
+
+mcp.tool(search4emdb, name="search_emdb")
+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 mmCIF file for given UniProt accessions.
+
+    Args:
+        uniprot_accs: A set of UniProt accessions.
+        save_dir: The directory to save the fetched files.
+
+    Returns:
+        A list of AlphaFold entries.
+    """
+    formats: set[DownloadableFormat] = {"cif"}
+    return alphafold_fetch(uniprot_accs, save_dir, formats)
+
+
+@mcp.tool
+async def fetch_emdb_volumes(emdb_ids: set[str], save_dir: Path) -> Mapping[str, Path]:
+    """Fetch EMDB volumes for given EMDB IDs.
+
+    Args:
+        emdb_ids: A set of EMDB IDs.
+        save_dir: The directory to save the fetched files.
+    Returns:
+        A mapping of EMDB ID to the path of the fetched volume file.
+    """
+    return await emdb_fetch(emdb_ids=emdb_ids, save_dir=save_dir)
+
+
+@mcp.tool
+def alphafold_confidence_filter(file: Path, query: ConfidenceFilterQuery, filtered_dir: Path) -> ConfidenceFilterResult:
+    """Take a mmcif/PDB file and filter it based on confidence (plDDT) scores.
+
+    If passes filter writes file to filtered_dir with residues above confidence threshold.
+    """
+    return filter_file_on_confidence(file, query, filtered_dir)
+
+
+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",
+    cellular_location: str = "nucleus",
+    confidence: int = 90,
+    min_residues: int = 100,
+    max_residues: int = 200,
+) -> str:
+    """Prompt to find candidate structures.
+
+    Args:
+        species: The species to search for (default: "Human").
+        cellular_location: The cellular location to search for (default: "nucleus").
+        confidence: The confidence threshold for AlphaFold structures (default: 90).
+        min_residues: Minimum number of high confidence residues (default: 100).
+        max_residues: Maximum number of high confidence residues (default: 200).
+
+    Returns:
+        A prompt string to find candidate structures.
+    """
+    return dedent(f"""\
+        Given the species '{species}' and cellular location '{cellular_location}' find the candidate structures.
+        Download structures from 2 sources namely PDB and Alphafold.
+        For alphafold I only want to use high confidence scores of over {confidence}.
+        and only keep structures with number of high confidence residues between {min_residues} and {max_residues}.
+
+        1. Search uniprot for proteins related to {species} and {cellular_location}.
+            1. For the species find the NCBI taxonomy id.
+            2. For cellular location find the associated GO term.
+            3. Find uniprot accessions based on NCBI taxonomy id and cellular location GO term.
+        2. For PDB
+            1. Search for structures related to the identified proteins.
+            2. Download each PDB entry from PDBe
+            3. Extract chain for the protein of interest.
+        3. For Alphafold
+            1. Search for AlphaFold entries related to the identified proteins.
+            2. Download each AlphaFold entry.
+            3. Filter the structures based on {confidence} as confidence
+               and nr residues between {min_residues} and {max_residues}.
+    """)