diff --git a/pixi.lock b/pixi.lock index 2ab955a..42f15b5 100644 --- a/pixi.lock +++ b/pixi.lock @@ -8069,6 +8069,8 @@ packages: - aiolimiter>=1.2,<2.0 - hishel>=1.3,<2.0 - loguru>=0.7.3,<1.0.0 + - pydantic>=2.13,<3.0 + - anysqlite>=0.0.5,<0.0.6 requires_python: '>=3.12' - pypi: https://files.pythonhosted.org/packages/0d/fe/6bea5c9162869c5beba5d9c8abbed835ec85bf1ec1fba05a3822325c45f3/build-1.5.0-py3-none-any.whl name: build diff --git a/py.typed b/py.typed new file mode 100644 index 0000000..e69de29 diff --git a/pyproject.toml b/pyproject.toml index d6bb02a..13171bf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -29,6 +29,8 @@ dependencies = [ "aiolimiter>=1.2,<2.0", "hishel>=1.3,<2.0", "loguru>=0.7.3,<1.0.0", + "pydantic>=2.13,<3.0", + "anysqlite>=0.0.5,<0.0.6", ] [project.urls] @@ -48,5 +50,5 @@ version_file = "reqadence/_version.py" version_scheme = "guess-next-dev" local_scheme = "node-and-timestamp" -[tool.basedpyright] -typeCheckingMode = "standard" \ No newline at end of file +[tool.setuptools.package-data] +"reqadence" = ["py.typed"] \ No newline at end of file diff --git a/reqadence/__init__.py b/reqadence/__init__.py index 6ec3575..1bc6af1 100644 --- a/reqadence/__init__.py +++ b/reqadence/__init__.py @@ -8,7 +8,7 @@ import sys from ast import literal_eval from importlib.metadata import PackageNotFoundError, version -from typing import Any +from typing import cast from loguru import logger @@ -40,9 +40,9 @@ def enable_logging( level: Requested log level: `10` is debug, `20` is info. file_path: Also write logs to files here. """ - config: dict[str, Any] = {"handlers": []} + handlers: list[dict[str, object]] = [] if stdout_set: - config["handlers"].append( + handlers.append( { "sink": sys.stdout, "level": level_set, @@ -51,7 +51,7 @@ def enable_logging( } ) if isinstance(file_path, str): - config["handlers"].append( + handlers.append( { "sink": file_path, "level": level_set, @@ -59,14 +59,13 @@ def enable_logging( "colorize": colorize, } ) + _ = logger.configure(handlers=handlers) # pyright: ignore[reportArgumentType] # https://loguru.readthedocs.io/en/stable/api/logger.html#loguru._logger.Logger.configure - logger.configure(**config) - logger.enable("reqadence") -if literal_eval(os.environ.get("REQADENCE_LOG", "False")): - level = int(os.environ.get("REQADENCE_LOG_LEVEL", 20)) - stdout = literal_eval(os.environ.get("REQADENCE_STDOUT", "True")) - log_file_path = os.environ.get("REQADENCE_LOG_FILE_PATH", None) +if cast(bool, literal_eval(os.environ.get("REQADENCE_LOG", "False"))): + level = int(os.environ.get("REQADENCE_LOG_LEVEL") or "20") + stdout = cast(bool, literal_eval(os.environ.get("REQADENCE_STDOUT") or "True")) + log_file_path = os.environ.get("REQADENCE_LOG_FILE_PATH") enable_logging(level, stdout, log_file_path) diff --git a/reqadence/api/base.py b/reqadence/api/base.py index e42ce3e..9909a11 100644 --- a/reqadence/api/base.py +++ b/reqadence/api/base.py @@ -13,7 +13,7 @@ from dataclasses import dataclass, field from enum import StrEnum from types import TracebackType -from typing import Self +from typing import Self, cast, override import httpx from aiolimiter import AsyncLimiter @@ -156,13 +156,16 @@ def _retry_delay(self, resp: httpx.Response | None, attempt: int) -> float: Returns: The delay in seconds before the next retry attempt. """ - if resp is not None and (value := resp.headers.get("retry-after")): + retry_after: str | None = ( + resp.headers.get("retry-after") if resp is not None else None + ) + if retry_after is not None: try: - return float(value) + return float(retry_after) except ValueError: logger.warning( "Invalid Retry-After header value: {value}. Falling back to backoff delay.", - value=value, + value=retry_after, ) return self.config.retry_policy.backoff_delay(attempt) @@ -266,7 +269,7 @@ async def _get_json( """GET and parse JSON, returning None on failure.""" resp = await self._get(url, params=params) try: - return resp.json() + return cast(JSONValue, resp.json()) except ValueError: logger.error("Failed to parse JSON from {url}", url=url) return None @@ -288,6 +291,7 @@ async def __aexit__( """Ensure the underlying http client is closed when exiting the async context manager.""" await self.aclose() + @override def __repr__(self) -> str: """Return the class and base URL for debugging purposes.""" return f"{self.__class__.__name__}(base_url={self.base_url!r})" diff --git a/reqadence/api/cache.py b/reqadence/api/cache.py index feb5517..55efe90 100644 --- a/reqadence/api/cache.py +++ b/reqadence/api/cache.py @@ -11,7 +11,7 @@ from abc import ABC, abstractmethod from collections.abc import Callable from dataclasses import dataclass -from typing import Any +from typing import Any, override import httpx from hishel import ( @@ -42,7 +42,7 @@ class CachePolicy(ABC): Default is .cache/hishel/hishel_cache.db""" @abstractmethod - def _build_hishel_policy(self) -> Any: + def _build_hishel_policy(self) -> SpecificationPolicy | FilterPolicy: """Return the hishel policy object for this strategy.""" raise NotImplementedError() @@ -66,13 +66,13 @@ def build_client_factory( default_ttl=self.default_ttl, ) - def factory(**kwargs: Any) -> httpx.AsyncClient: + def factory(**kwargs: Any) -> httpx.AsyncClient: # pyright: ignore[reportAny, reportExplicitAny] transport = AsyncCacheTransport( next_transport=next_transport or httpx.AsyncHTTPTransport(), storage=storage, policy=cache_policy, ) - return httpx.AsyncClient(transport=transport, **kwargs) + return httpx.AsyncClient(transport=transport, **kwargs) # pyright: ignore[reportAny] return factory @@ -85,7 +85,8 @@ class RFCCachePolicy(CachePolicy): Use for APIs that send proper cache headers. """ - def _build_hishel_policy(self) -> Any: + @override + def _build_hishel_policy(self) -> SpecificationPolicy: return SpecificationPolicy( cache_options=CacheOptions(supported_methods=list(self.supported_methods)) @@ -98,25 +99,30 @@ class AlwaysCachePolicy(CachePolicy): cacheable_status_codes: frozenset[int] = frozenset({200}) - def _build_hishel_policy(self) -> Any: + @override + def _build_hishel_policy(self) -> FilterPolicy: methods = frozenset(m.upper() for m in self.supported_methods) codes = frozenset(self.cacheable_status_codes) class _MethodFilter(BaseFilter[Request]): + @override def needs_body(self) -> bool: """This filter does not require the request body.""" return False + @override def apply(self, item: Request, body: bytes | None) -> bool: """Cache only if the HTTP method is in the supported methods set.""" return item.method.upper() in methods class _StatusFilter(BaseFilter[Response]): + @override def needs_body(self) -> bool: """This filter does not require the response body.""" return False + @override def apply(self, item: Response, body: bytes | None) -> bool: """ Cache only if the HTTP status code is in the cacheable diff --git a/reqadence/api/errors.py b/reqadence/api/errors.py index 0c4bd2a..2762021 100644 --- a/reqadence/api/errors.py +++ b/reqadence/api/errors.py @@ -9,9 +9,9 @@ class APIError(Exception): def __init__(self, url: str, status_code: int | None = None) -> None: """Base class for API failures, distinguishing retryable errors from permanent ones.""" - self.url = url + self.url: str = url """The URL of the API endpoint.""" - self.status_code = status_code + self.status_code: int | None = status_code """The HTTP status code of the API response, if available.""" message = f"{url} (HTTP {status_code})" if status_code else url super().__init__(message) diff --git a/reqadence/api/rcsb/client.py b/reqadence/api/rcsb/client.py index 35488f4..05d37ba 100644 --- a/reqadence/api/rcsb/client.py +++ b/reqadence/api/rcsb/client.py @@ -4,16 +4,24 @@ """RCSBClient Protein Data Bank API client and entry model.""" -from typing import Any, Callable +import asyncio +from collections.abc import Awaitable, Callable +from typing import TypeVar import httpx +from aiolimiter import AsyncLimiter from loguru import logger -from pydantic import ValidationError +from pydantic import BaseModel, ValidationError -from reqadence.api.base import BaseAPI, JSONValue +from reqadence.api.base import BaseAPI, ClientConfig from reqadence.api.cache import AlwaysCachePolicy from reqadence.api.errors import PermanentAPIError -from reqadence.api.rcsb.model import RCSBEntry +from reqadence.api.rcsb.model import ( + ChemComp, + NonpolymerEntity, + PolymerEntity, + RCSBEntry, +) RCSB_BASE_URL = "https://data.rcsb.org/rest/v1" """Base URL for the RCSBClient REST API endpoints.""" @@ -21,6 +29,8 @@ RCSB_FILES_URL = "https://files.rcsb.org/download" """Base URL for the RCSBClient file download service.""" +_ModelT = TypeVar("_ModelT", bound=BaseModel) + class RCSBClient(BaseAPI): """Client for the RCSBClient REST API and file download service. @@ -31,7 +41,10 @@ def __init__( self, base_url: str = RCSB_BASE_URL, client_factory: Callable[..., httpx.AsyncClient] | None = None, - **kwargs: Any, + *, + config: ClientConfig | None = None, + rate_limiter: AsyncLimiter | None = None, + sleep: Callable[[float], Awaitable[None]] = asyncio.sleep, ) -> None: """Initialize the RCSBClient with caching enabled by default. @@ -41,7 +54,13 @@ def __init__( """ if client_factory is None: client_factory = AlwaysCachePolicy().build_client_factory() - super().__init__(base_url=base_url, client_factory=client_factory, **kwargs) + super().__init__( + base_url=base_url, + client_factory=client_factory, + config=config, + rate_limiter=rate_limiter, + sleep=sleep, + ) @staticmethod def _to_legacy(pdb_id: str) -> str | None: @@ -57,41 +76,52 @@ def _to_legacy(pdb_id: str) -> str | None: return pdb_id[-4:] return None - async def get_entry(self, pdb_id: str) -> RCSBEntry | None: - """Fetch an entry record. Accepts legacy or extended PDB IDs. + async def _get_model( + self, url: str, model: type[_ModelT], context: str + ) -> _ModelT | None: + """GET a JSON object and validate it into the given pydantic model. Args: - pdb_id: 4-character legacy PDB ID or 12-character extended ID. + url: Endpoint path to fetch. + model: The pydantic model class to validate the payload into. + context: Human-readable subject for log messages (e.g. the PDB ID). Returns: - An `RCSBEntry` object if found, or None if not found. + The validated model instance, or None if the request failed, the + payload was not a JSON object, or validation failed. """ - - pdb_id = self._to_legacy(pdb_id) or pdb_id try: - data = await self._get_json(f"core/entry/{pdb_id}") + data = await self._get_json(url) except PermanentAPIError: return None if not isinstance(data, dict): - logger.error( - "Unexpected response type for {pdb_id}: {type(data)}", - pdb_id=pdb_id, - type=type(data).__name__, - ) return None try: - return RCSBEntry.model_validate(data) + return model.model_validate(data) except ValidationError as exc: logger.error( - "Failed to validate RCSBEntry for {pdb_id}: {exc}", - pdb_id=pdb_id, + "Failed to validate {model} for {context}: {exc}", + model=model.__name__, + context=context, exc=exc, ) return None + async def get_entry(self, pdb_id: str) -> RCSBEntry | None: + """Fetch an entry record. Accepts legacy or extended PDB IDs. + + Args: + pdb_id: 4-character legacy PDB ID or 12-character extended ID. + + Returns: + An `RCSBEntry` object if found, or None if not found. + """ + pdb_id = self._to_legacy(pdb_id) or pdb_id + return await self._get_model(f"core/entry/{pdb_id}", RCSBEntry, context=pdb_id) + async def _get_polymer_entity( self, pdb_id: str, entity_id: str - ) -> dict[str, Any] | None: + ) -> PolymerEntity | None: """Fetch a polymer entity record for the given PDB ID and entity ID. Args: @@ -99,18 +129,18 @@ async def _get_polymer_entity( entity_id: The polymer entity ID. Returns: - A dictionary containing the polymer entity data, or None if not found. + A `PolymerEntity` if found, or None if not found. """ pdb_id = self._to_legacy(pdb_id) or pdb_id - try: - data = await self._get_json(url=f"core/polymer_entity/{pdb_id}/{entity_id}") - except PermanentAPIError: - return None - return data if isinstance(data, dict) else None + return await self._get_model( + f"core/polymer_entity/{pdb_id}/{entity_id}", + PolymerEntity, + context=f"{pdb_id}/{entity_id}", + ) async def _get_nonpolymer_entity( self, pdb_id: str, entity_id: str - ) -> dict[str, Any] | None: + ) -> NonpolymerEntity | None: """Fetch a non-polymer entity record for the given PDB ID and entity ID. Args: @@ -118,30 +148,27 @@ async def _get_nonpolymer_entity( entity_id: The non-polymer entity ID. Returns: - A dictionary containing the non-polymer entity data, or None if not found.""" + A `NonpolymerEntity` if found, or None if not found. + """ pdb_id = self._to_legacy(pdb_id) or pdb_id - try: - data = await self._get_json( - url=f"core/nonpolymer_entity/{pdb_id}/{entity_id}" - ) - except PermanentAPIError: - return None - return data if isinstance(data, dict) else None + return await self._get_model( + f"core/nonpolymer_entity/{pdb_id}/{entity_id}", + NonpolymerEntity, + context=f"{pdb_id}/{entity_id}", + ) - async def _get_chemcomp(self, comp_id: str) -> dict[str, Any] | None: + async def _get_chemcomp(self, comp_id: str) -> ChemComp | None: """Fetch a chemical component record (e.g. for SMILES lookup). Args: comp_id: The chemical component ID. Returns: - A dictionary containing the chemical component data, or None if not found. + A `ChemComp` if found, or None if not found. """ - try: - data = await self._get_json(f"core/chemcomp/{comp_id}") - except PermanentAPIError: - return None - return data if isinstance(data, dict) else None + return await self._get_model( + f"core/chemcomp/{comp_id}", ChemComp, context=comp_id + ) async def get_structure(self, pdb_id: str) -> tuple[str, str] | None: """Download a structure as PDB, falling back to mmCIF if PDB is not available. @@ -186,27 +213,26 @@ async def get_ligand_smiles(self, comp_id: str) -> dict[str, str]: return self.extract_smiles(data) @staticmethod - def extract_smiles(chemcomp_data: dict[str, Any]) -> dict[str, str]: + def extract_smiles(chemcomp: ChemComp) -> dict[str, str]: """Extract canonical and stereo SMILES from a chem component record. Falls back to pdbx_chem_comp_descriptor: SMILES_CANONICAL when the rcsb_chem_comp_descriptor block is missing the smiles field. Args: - chemcomp_data: The chemical component data dictionary from the API. + chemcomp: The parsed chemical component record from the API. Returns: A dictionary containing the SMILES and stereo SMILES, or an empty dictionary if not found """ - descriptor = chemcomp_data.get("rcsb_chem_comp_descriptor", {}) - if "SMILES" in descriptor: + descriptor = chemcomp.rcsb_chem_comp_descriptor + if descriptor is not None and descriptor.smiles is not None: return { - "smiles": descriptor["SMILES"], - "stereo_smiles": descriptor.get("SMILES_stereo", descriptor["SMILES"]), + "smiles": descriptor.smiles, + "stereo_smiles": descriptor.smiles_stereo or descriptor.smiles, } - for d in chemcomp_data.get("pdbx_chem_comp_descriptor", []): - if d.get("type") == "SMILES_CANONICAL": - smi = d.get("descriptor", "") - return {"smiles": smi, "stereo_smiles": smi} + for row in chemcomp.pdbx_chem_comp_descriptor: + if row.type == "SMILES_CANONICAL": + return {"smiles": row.descriptor, "stereo_smiles": row.descriptor} logger.warning("No SMILES found in chem component record.") return {"smiles": "", "stereo_smiles": ""} @@ -220,7 +246,7 @@ async def mutations_count(self, entry: RCSBEntry) -> dict[str, int]: Returns: A dictionary mapping polymer entity IDs to their mutation counts. """ - counts = {} + counts: dict[str, int] = {} for entity_id in entry.entity_ids["polymer"]: data = await self._get_polymer_entity(entry.pdb_id, entity_id) if data is None: @@ -230,8 +256,7 @@ async def mutations_count(self, entry: RCSBEntry) -> dict[str, int]: pdb_id=entry.pdb_id, ) continue - count = data.get("entity_poly", {}).get("rcsb_mutation_count", 0) - counts[entity_id] = count + counts[entity_id] = data.entity_poly.rcsb_mutation_count return counts async def nonpolymer_names(self, entry: RCSBEntry) -> dict[str, str]: @@ -245,7 +270,7 @@ async def nonpolymer_names(self, entry: RCSBEntry) -> dict[str, str]: """ if not entry.has_ligands: return {} - names = {} + names: dict[str, str] = {} for entity_id in entry.entity_ids["nonpolymer"]: data = await self._get_nonpolymer_entity(entry.pdb_id, entity_id) if data is None: @@ -255,7 +280,7 @@ async def nonpolymer_names(self, entry: RCSBEntry) -> dict[str, str]: pdb_id=entry.pdb_id, ) continue - comp_id = data.get("pdbx_entity_nonpoly", {}).get("comp_id") + comp_id = data.pdbx_entity_nonpoly.comp_id if comp_id: names[entity_id] = comp_id return names diff --git a/reqadence/api/rcsb/model.py b/reqadence/api/rcsb/model.py index 7eac0a9..2bf77ee 100644 --- a/reqadence/api/rcsb/model.py +++ b/reqadence/api/rcsb/model.py @@ -176,3 +176,79 @@ def entity_ids(self) -> dict[str, list[str]]: "polymer": ids.polymer_entity_ids, "nonpolymer": ids.non_polymer_entity_ids, } + + +class _ChemCompDescriptor(BaseModel): + """Primary SMILES descriptor block from a chemical component record.""" + + smiles: str | None = Field( + default=None, + alias="SMILES", + description="Canonical SMILES string for the chemical component.", + ) + smiles_stereo: str | None = Field( + default=None, + alias="SMILES_stereo", + description="Stereochemistry-aware SMILES, if reported.", + ) + + +class _PdbxChemCompDescriptor(BaseModel): + """One row of the fallback pdbx_chem_comp_descriptor list.""" + + type: str = "" + """Descriptor type (e.g. ``SMILES_CANONICAL``).""" + + descriptor: str = "" + """The descriptor string itself (e.g. a SMILES when ``type`` is a SMILES variant).""" + + +class ChemComp(BaseModel): + """A parsed RCSB chemical component record.""" + + model_config: ClassVar[ConfigDict] = ConfigDict(populate_by_name=True) + + rcsb_chem_comp_descriptor: _ChemCompDescriptor | None = Field( + default=None, + description="Primary descriptor block; preferred SMILES source.", + ) + pdbx_chem_comp_descriptor: list[_PdbxChemCompDescriptor] = Field( + default_factory=list, + description="Fallback descriptor rows used when the primary block lacks SMILES.", + ) + + +class _EntityPoly(BaseModel): + """Polymer-entity sub-block carrying the mutation count.""" + + rcsb_mutation_count: int = 0 + """Number of reported mutations for this polymer entity.""" + + +class PolymerEntity(BaseModel): + """A parsed RCSB polymer-entity record.""" + + model_config: ClassVar[ConfigDict] = ConfigDict(populate_by_name=True) + + entity_poly: _EntityPoly = Field( + default_factory=_EntityPoly, + description="Polymer-entity block holding the mutation count.", + ) + + +class _PdbxEntityNonpoly(BaseModel): + """Non-polymer-entity sub-block carrying the component id.""" + + comp_id: str | None = None + """Chemical-component ID of the non-polymer entity, if reported.""" + + +class NonpolymerEntity(BaseModel): + """A parsed RCSB non-polymer-entity record.""" + + model_config: ClassVar[ConfigDict] = ConfigDict(populate_by_name=True) + + pdbx_entity_nonpoly: _PdbxEntityNonpoly = Field( + default_factory=_PdbxEntityNonpoly, + description="Non-polymer-entity block holding the component id.", + ) diff --git a/tests/test_rcsb.py b/tests/test_rcsb.py index c645b50..54a7f96 100644 --- a/tests/test_rcsb.py +++ b/tests/test_rcsb.py @@ -15,6 +15,7 @@ from reqadence.api.base import ClientConfig from reqadence.api.cache import AlwaysCachePolicy, RFCCachePolicy from reqadence.api.rcsb import RCSBClient, RCSBEntry +from reqadence.api.rcsb.model import ChemComp from reqadence.api.retry import RetryPolicy _ENTRY_6OAV = { @@ -244,7 +245,7 @@ def test_extract_smiles(): "SMILES_stereo": "CC[OH]", } } - assert RCSBClient.extract_smiles(data) == { + assert RCSBClient.extract_smiles(ChemComp.model_validate(data)) == { "smiles": "CCO", "stereo_smiles": "CC[OH]", } @@ -256,12 +257,12 @@ def test_extract_smiles(): {"type": "SMILES_CANONICAL", "descriptor": "CCO"}, ], } - assert RCSBClient.extract_smiles(data) == { + assert RCSBClient.extract_smiles(ChemComp.model_validate(data)) == { "smiles": "CCO", "stereo_smiles": "CCO", } - assert RCSBClient.extract_smiles({}) == {"smiles": "", "stereo_smiles": ""} + assert RCSBClient.extract_smiles(ChemComp()) == {"smiles": "", "stereo_smiles": ""} async def test_get_entry(rcsb):