Source code for boltz_data.ccd._from_rcsb

"""Fetch chemical components from RCSB PDB."""

from collections.abc import Iterator, Mapping
from functools import cache

from boltz_data.cif._read import read_single_cif_from_file

from ._from_mmcif import read_chemical_component_dictionary_from_mmcif
from ._models import ChemicalComponent




def chemical_component_from_rcsb(*, comp_id: str) -> ChemicalComponent:
    """
    Fetch a chemical component from RCSB PDB.

    Args:
        comp_id: The component ID to fetch.

    Returns:
        The chemical component.

    Raises:
        ValueError: If the component ID is not found in RCSB PDB.

    """
    url = f"https://files.rcsb.org/ligands/download/{comp_id}.cif"
    mmcif = read_single_cif_from_file(url)
    chemical_components = read_chemical_component_dictionary_from_mmcif(mmcif)
    if comp_id not in chemical_components:
        msg = f"Component ID {comp_id} not found in RCSB PDB."
        raise ValueError(msg)
    return chemical_components[comp_id]



class RemoteChemicalComponentDatabase(Mapping[str, ChemicalComponent]):
    """A mapping of chemical component IDs to ChemicalComponent objects fetched from RCSB PDB on demand."""

    def __init__(self) -> None:
        """Initialize the remote database with an empty cache."""
        self._cache: dict[str, ChemicalComponent] = {}

    def __getitem__(self, key: str) -> ChemicalComponent:
        """Fetch a chemical component, using cache if available."""
        if key not in self._cache:
            self._cache[key] = chemical_component_from_rcsb(comp_id=key)
        return self._cache[key]

    def __iter__(self) -> Iterator[str]:
        """Raise NotImplementedError as iteration is not supported."""
        msg = "Iteration is not supported for RemoteChemicalComponentDatabase."
        raise NotImplementedError(msg)

    def __len__(self) -> int:
        """Raise NotImplementedError as length is not supported."""
        msg = "Length is not supported for RemoteChemicalComponentDatabase."
        raise NotImplementedError(msg)




@cache
def get_remote_chemical_component_database() -> RemoteChemicalComponentDatabase:
    """Get a chemical component database that talks directly to the RCSB API."""
    return RemoteChemicalComponentDatabase()