from collections.abc import Iterator
from .data import REFMOL_DICT, RefMolID
from .dclass import RefMolInfo
[docs]
class RefMolLookup:
def __init__(self, refmol_data: dict[RefMolID, RefMolInfo]) -> None:
self._refmol_data = refmol_data
self._refmolid_to_data: dict[RefMolID, RefMolInfo] = {}
self._name_to_data: dict[str, RefMolInfo] = {}
self._label_type_to_data: dict[str, list[RefMolInfo]] = {}
self._molecule_type_to_data: dict[str, list[RefMolInfo]] = {}
for refmol_id, refmol_info in refmol_data.items():
self._refmolid_to_data[refmol_id] = refmol_info
self._name_to_data[refmol_info.name.lower()] = refmol_info
# Group by label type
if refmol_info.label_type:
self._label_type_to_data.setdefault(refmol_info.label_type.lower(), []).append(refmol_info)
# Group by molecule type
if refmol_info.molecule_type:
self._molecule_type_to_data.setdefault(refmol_info.molecule_type.lower(), []).append(refmol_info)
[docs]
def query_id(self, refmol_id: RefMolID) -> RefMolInfo | None:
"""Query by RefMolID enum"""
return self._refmolid_to_data.get(refmol_id)
[docs]
def query_name(self, name: str) -> RefMolInfo | None:
"""Query by reference molecule name (e.g., 'TMT126', 'sidechain_A')"""
return self._name_to_data.get(name.lower())
[docs]
def query_label_type(self, label_type: str) -> list[RefMolInfo]:
"""Query all molecules by label type (e.g., 'TMT', 'iTRAQ')"""
return self._label_type_to_data.get(label_type.lower(), [])
[docs]
def query_molecule_type(self, molecule_type: str) -> list[RefMolInfo]:
"""Query all molecules by molecule type (e.g., 'reporter', 'sidechain', 'nucleobase')"""
return self._molecule_type_to_data.get(molecule_type.lower(), [])
def __getitem__(self, key: str | RefMolID) -> RefMolInfo:
"""Get reference molecule by ID or name"""
if isinstance(key, RefMolID):
info = self.query_id(key)
if info is not None:
return info
raise KeyError(f"Reference molecule ID '{key}' not found.")
# Try name
info = self.query_name(key)
if info is not None:
return info
raise KeyError(f"Reference molecule '{key}' not found by name.")
def __contains__(self, key: str | RefMolID) -> bool:
"""Check if reference molecule exists"""
try:
self[key]
return True
except KeyError:
return False
[docs]
def get(self, key: str | RefMolID) -> RefMolInfo | None:
"""Get reference molecule or None if not found"""
try:
return self[key]
except KeyError:
return None
def __iter__(self) -> Iterator[RefMolInfo]:
"""Iterator over all RefMolInfo entries in the lookup."""
return iter(self._refmol_data.values())
REFMOL_LOOKUP = RefMolLookup(REFMOL_DICT)