|
|
|
|
| """Library for loading structure data from various sources.""" |
|
|
| from collections.abc import Mapping, Sequence |
| import functools |
| import os |
| import pathlib |
| import tarfile |
|
|
|
|
| class NotFoundError(KeyError): |
| """Raised when the structure store doesn't contain the requested target.""" |
|
|
|
|
| class StructureStore: |
| """Handles the retrieval of mmCIF files from a filesystem.""" |
|
|
| def __init__( |
| self, |
| structures: str | os.PathLike[str] | Mapping[str, str], |
| ): |
| """Initialises the instance. |
| |
| Args: |
| structures: Path of the directory where the mmCIF files are or a Mapping |
| from target name to mmCIF string. |
| """ |
| if isinstance(structures, Mapping): |
| self._structure_mapping = structures |
| self._structure_path = None |
| self._structure_tar = None |
| else: |
| self._structure_mapping = None |
| path_str = os.fspath(structures) |
| if path_str.endswith('.tar'): |
| self._structure_tar = tarfile.open(path_str, 'r') |
| self._structure_path = None |
| else: |
| self._structure_path = pathlib.Path(structures) |
| self._structure_tar = None |
|
|
| @functools.cached_property |
| def _tar_members(self) -> Mapping[str, tarfile.TarInfo]: |
| assert self._structure_tar is not None |
| return { |
| path.stem: tarinfo |
| for tarinfo in self._structure_tar.getmembers() |
| if tarinfo.isfile() |
| and (path := pathlib.Path(tarinfo.path.lower())).suffix == '.cif' |
| } |
|
|
| def get_mmcif_str(self, target_name: str) -> str: |
| """Returns an mmCIF for a given `target_name`. |
| |
| Args: |
| target_name: Name specifying the target mmCIF. |
| |
| Raises: |
| NotFoundError: If the target is not found. |
| """ |
| if self._structure_mapping is not None: |
| try: |
| return self._structure_mapping[target_name] |
| except KeyError as e: |
| raise NotFoundError(f'{target_name=} not found') from e |
|
|
| if self._structure_tar is not None: |
| try: |
| member = self._tar_members[target_name] |
| if struct_file := self._structure_tar.extractfile(member): |
| return struct_file.read().decode() |
| else: |
| raise NotFoundError(f'{target_name=} not found') |
| except KeyError: |
| raise NotFoundError(f'{target_name=} not found') from None |
|
|
| filepath = self._structure_path / f'{target_name}.cif' |
| try: |
| return filepath.read_text() |
| except FileNotFoundError as e: |
| raise NotFoundError(f'{target_name=} not found at {filepath=}') from e |
|
|
| def target_names(self) -> Sequence[str]: |
| """Returns all targets in the store.""" |
| if self._structure_mapping is not None: |
| return [*self._structure_mapping.keys()] |
| elif self._structure_tar is not None: |
| return sorted(self._tar_members.keys()) |
| elif self._structure_path is not None: |
| return sorted([path.stem for path in self._structure_path.glob('*.cif')]) |
| return () |
|
|
