"""Core dataset loading implementation."""
from __future__ import annotations
import itertools
import json
import logging
from collections.abc import Iterator, Sequence
from pathlib import Path
from urllib.parse import urlencode
import polars as pl
from pydantic import ValidationError
from ._cache import DownloadCache
from ._filters import MoleculeFilters
from ._http import fetch_json
from .dataset._molecules import MoleculeDataset
from .dataset._source import EagerSource
from .exceptions import (
DatasetNotFoundError,
DownloadError,
ResultTooLargeError,
SchemaError,
)
from .settings import settings
from .types import (
DataFormat,
Include,
MoleculePage,
SubsetData,
SubsetMetadata,
Value,
)
logger = logging.getLogger(__name__)
_MOLECULE_PAGE_SIZE = 100
def _load_dataset_impl(
subset: str,
*,
include: Sequence[Include] = ("molecules", "reactions", "fragments"),
data_format: DataFormat | str = DataFormat.PARQUET,
include_structures: bool = False,
molecule_properties: Sequence[str] | None = None,
reaction_properties: Sequence[str] | None = None,
fragment_properties: Sequence[str] | None = None,
smarts: str | None = None,
cache_dir: str | Path | None = None,
force_download: bool = False,
) -> SubsetData:
"""Download (or load from cache) a subset and return its data.
Args:
subset: Subset slug to download.
include: Entity types to include in the download.
data_format: Serialisation format for entity data files.
include_structures: Whether to download XYZ structure files.
molecule_properties: Restrict molecule columns to these properties.
reaction_properties: Restrict reaction columns to these properties.
fragment_properties: Restrict fragment columns to these properties.
smarts: Filter downloaded molecules by SMARTS substructure.
cache_dir: Override the default cache directory.
force_download: Re-download even if cached data exists.
"""
fmt = DataFormat(data_format) if isinstance(data_format, str) else data_format
include_set = _build_include_set(include, include_structures)
smarts = smarts.strip() if smarts else None
url = _build_download_url(
subset,
include_set,
fmt,
molecule_properties,
reaction_properties,
fragment_properties,
smarts=smarts,
)
cache = DownloadCache(Path(cache_dir) if cache_dir is not None else None)
cache_key = DownloadCache.compute_cache_key(
subset,
sorted(include_set),
fmt.value,
molecule_properties,
reaction_properties,
fragment_properties,
smarts=smarts,
)
data_dir = _download(cache, url, subset, cache_key, force_download)
result = SubsetData._from_download(
data_dir=data_dir,
data_format=fmt,
molecules=_read_entity(data_dir, "molecules", fmt)
if "molecules" in include_set
else None,
reactions=_read_entity(data_dir, "reactions", fmt)
if "reactions" in include_set
else None,
fragments=_read_entity(data_dir, "fragments", fmt)
if "fragments" in include_set
else None,
metadata=_read_metadata(data_dir),
)
if "structures" in include_set and result.molecules is not None:
result.molecules = _attach_structure_paths(
result.molecules, data_dir / "structures"
)
_trigger_stub_sync(cache)
return result
[docs]
def load_molecules(
*,
subset: str | None = None,
smarts: str | None = None,
smiles: str | None = None,
molecular_formula: str | None = None,
limit: int | None = None,
) -> MoleculeDataset[dict[str, Value]]:
"""Load molecules from the list endpoint with simple server-side filters.
Filters are optional, single-value, and combined with ``AND`` semantics.
Pass ``subset`` to restrict the search to one subset; omit it to search
across all molecules in the database.
Args:
subset: Restrict the search to a single subset slug.
smarts: SMARTS substructure pattern.
smiles: SMILES substring.
molecular_formula: Molecular-formula substring.
limit: Stop after this many molecules. Passing a ``limit`` also opts out
of the unbounded-result safety check.
Raises:
ResultTooLargeError: When no ``limit`` is given and the query matches
more than ``settings.max_molecules`` molecules. Narrow the query,
pass ``limit``, or raise ``LCMD_DB_MAX_MOLECULES``.
"""
if limit is not None and limit < 0:
raise ValueError("limit must be non-negative")
filters = MoleculeFilters(
smarts=smarts,
smiles=smiles,
molecular_formula=molecular_formula,
)
url = _build_molecule_list_url(subset=subset, filters=filters)
first_page = _fetch_molecule_page(url)
_guard_molecule_count(first_page.count, limit=limit)
rows: Iterator[dict[str, object]] = _iter_molecule_rows(first_page)
if limit is not None:
rows = itertools.islice(rows, limit)
return MoleculeDataset(EagerSource(pl.DataFrame(list(rows))))
def _guard_molecule_count(count: int, *, limit: int | None) -> None:
"""Refuse to stream an unbounded full-table result above the safety ceiling."""
ceiling = settings.max_molecules
if limit is not None or ceiling <= 0:
return
if count > ceiling:
raise ResultTooLargeError(
f"load_molecules matched {count} molecules, above the safety ceiling "
f"of {ceiling}. Narrow the query, pass limit=, or raise the ceiling "
f"via the LCMD_DB_MAX_MOLECULES environment variable."
)
[docs]
def clear_cache(
subset: str | None = None,
cache_dir: str | Path | None = None,
) -> None:
"""Remove cached data files.
Args:
subset: If given, only clear the cache for this subset. Otherwise clear all.
cache_dir: Override the default cache directory.
"""
DownloadCache(Path(cache_dir) if cache_dir is not None else None).clear(subset)
_VALID_INCLUDE: frozenset[str] = frozenset(
{"molecules", "reactions", "fragments", "structures"}
)
def _build_include_set(include: Sequence[str], include_structures: bool) -> set[str]:
invalid = set(include) - _VALID_INCLUDE
if invalid:
raise ValueError(
f"Invalid include values: {sorted(invalid)}. "
f"Valid options: {sorted(_VALID_INCLUDE)}"
)
result = set(include)
if include_structures:
result |= {"structures", "molecules"}
return result
def _build_download_url(
subset: str,
include_set: set[str],
fmt: DataFormat,
molecule_properties: Sequence[str] | None,
reaction_properties: Sequence[str] | None,
fragment_properties: Sequence[str] | None,
smarts: str | None = None,
) -> str:
params: dict[str, str] = {
"include": ",".join(sorted(include_set)),
"data_format": fmt.value,
}
for key, props in [
("molecule_properties", molecule_properties),
("reaction_properties", reaction_properties),
("fragment_properties", fragment_properties),
]:
if props is not None:
params[key] = ",".join(props)
if smarts:
params["smarts"] = smarts
return f"{settings.base_url}/api/v1/subsets/{subset}/download/?{urlencode(params)}"
def _build_molecule_list_url(
*,
subset: str | None,
filters: MoleculeFilters,
page_size: int = _MOLECULE_PAGE_SIZE,
page: int | None = None,
) -> str:
params = filters.to_query_params()
if subset:
params["subset"] = subset
params["page_size"] = str(page_size)
if page is not None:
params["page"] = str(page)
return f"{settings.base_url}/api/v1/molecules/?{urlencode(params)}"
def _fetch_molecule_page(url: str) -> MoleculePage:
try:
return MoleculePage.model_validate(fetch_json(url))
except ValidationError as exc:
raise SchemaError(f"Unexpected molecule list response from {url}") from exc
def _iter_molecule_rows(first_page: MoleculePage) -> Iterator[dict[str, object]]:
page = first_page
while True:
yield from page.results
next_url = page.next
if not next_url:
return
page = _fetch_molecule_page(next_url)
def _download(
cache: DownloadCache,
url: str,
subset: str,
cache_key: str,
force: bool,
) -> Path:
try:
return cache.get_or_download(url, subset, cache_key, force=force)
except FileNotFoundError as exc:
raise DatasetNotFoundError(str(exc)) from exc
except Exception as exc:
if _is_not_found(exc):
raise DatasetNotFoundError(f"Subset '{subset}' not found") from exc
raise DownloadError(f"Failed to download dataset: {exc}") from exc
def _is_not_found(exc: BaseException) -> bool:
"""Check if the exception represents an HTTP 404."""
response = getattr(exc, "response", None)
if response is not None and getattr(response, "status_code", None) == 404:
return True
code = getattr(exc, "code", None)
if code == 404:
return True
return False
_READERS = {
DataFormat.PARQUET: pl.read_parquet,
DataFormat.CSV: pl.read_csv,
DataFormat.TSV: lambda p: pl.read_csv(p, separator="\t"),
DataFormat.XLSX: pl.read_excel,
DataFormat.JSON: pl.read_json,
}
def _read_entity(data_dir: Path, entity: str, fmt: DataFormat) -> pl.DataFrame | None:
data_file = data_dir / f"{entity}.{fmt.value}"
return _READERS[fmt](data_file) if data_file.exists() else None
def _attach_structure_paths(df: pl.DataFrame, structures_dir: Path) -> pl.DataFrame:
if not structures_dir.exists():
return df
return df.with_columns(
pl.col("id")
.map_elements(
lambda mid: str(structures_dir / f"{mid}.xyz"), return_dtype=pl.Utf8
)
.alias("structure_path")
)
def _read_metadata(data_dir: Path) -> SubsetMetadata | None:
meta_file = data_dir / "metadata.json"
if not meta_file.exists():
return None
return SubsetMetadata.model_validate(json.loads(meta_file.read_bytes()))
def _trigger_stub_sync(cache: DownloadCache) -> None:
if not settings.auto_sync_stubs:
return
from ._sync import sync_stubs
sync_stubs(cache)
![lcmd[db] - Home](../../_static/logo.svg){kind=logo}
![lcmd[db] - Home](../../_static/logo-dark.svg){kind=logo}