protein-quest 0.3.2__py3-none-any.whl → 0.5.0__py3-none-any.whl

protein_quest/uniprot.py

@@ -201,7 +201,7 @@ def _build_sparql_generic_query(select_clause: str, where_clause: str, limit: in
     """)
 
 
-def _build_sparql_generic_by_uniprot_accesions_query(
+def _build_sparql_generic_by_uniprot_accessions_query(
     uniprot_accs: Iterable[str], select_clause: str, where_clause: str, limit: int = 10_000, groupby_clause=""
 ) -> str:
     values = " ".join(f'("{ac}")' for ac in uniprot_accs)
@@ -269,7 +269,7 @@ def _build_sparql_query_pdb(uniprot_accs: Iterable[str], limit=10_000) -> str:
     """)
 
     groupby_clause = "?protein ?pdb_db ?pdb_method ?pdb_resolution"
-    return _build_sparql_generic_by_uniprot_accesions_query(
+    return _build_sparql_generic_by_uniprot_accessions_query(
         uniprot_accs, select_clause, where_clause, limit, groupby_clause
     )
 
@@ -284,7 +284,7 @@ def _build_sparql_query_af(uniprot_accs: Iterable[str], limit=10_000) -> str:
         ?protein rdfs:seeAlso ?af_db .
         ?af_db up:database <http://purl.uniprot.org/database/AlphaFoldDB> .
     """)
-    return _build_sparql_generic_by_uniprot_accesions_query(uniprot_accs, select_clause, dedent(where_clause), limit)
+    return _build_sparql_generic_by_uniprot_accessions_query(uniprot_accs, select_clause, dedent(where_clause), limit)
 
 
 def _build_sparql_query_emdb(uniprot_accs: Iterable[str], limit=10_000) -> str:
@@ -297,7 +297,7 @@ def _build_sparql_query_emdb(uniprot_accs: Iterable[str], limit=10_000) -> str:
         ?protein rdfs:seeAlso ?emdb_db .
         ?emdb_db up:database <http://purl.uniprot.org/database/EMDB> .
     """)
-    return _build_sparql_generic_by_uniprot_accesions_query(uniprot_accs, select_clause, dedent(where_clause), limit)
+    return _build_sparql_generic_by_uniprot_accessions_query(uniprot_accs, select_clause, dedent(where_clause), limit)
 
 
 def _execute_sparql_search(
@@ -509,3 +509,156 @@ def search4emdb(uniprot_accs: Iterable[str], limit: int = 10_000, timeout: int =
     )
     limit_check("Search for EMDB entries on uniprot", limit, len(raw_results))
     return _flatten_results_emdb(raw_results)
+
+
+def _build_complex_sparql_query(uniprot_accs: Iterable[str], limit: int) -> str:
+    """Builds a SPARQL query to retrieve ComplexPortal information for given UniProt accessions.
+
+    Example:
+
+    ```sparql
+    PREFIX up:   <http://purl.uniprot.org/core/>
+    PREFIX rdf:  <http://www.w3.org/1999/02/22-rdf-syntax-ns#>
+    PREFIX rdfs: <http://www.w3.org/2000/01/rdf-schema#>
+
+    SELECT
+    ?protein
+    ?cp_db
+    ?cp_comment
+    (GROUP_CONCAT(DISTINCT ?member; separator=",") AS ?complex_members)
+    (COUNT(DISTINCT ?member) AS ?member_count)
+    WHERE {
+    # Input UniProt accessions
+    VALUES (?ac) { ("P05067") ("P60709") ("Q05471")}
+    BIND (IRI(CONCAT("http://purl.uniprot.org/uniprot/", ?ac)) AS ?protein)
+
+    # ComplexPortal cross-reference for each input protein
+    ?protein a up:Protein ;
+            rdfs:seeAlso ?cp_db .
+    ?cp_db up:database <http://purl.uniprot.org/database/ComplexPortal> .
+    OPTIONAL { ?cp_db rdfs:comment ?cp_comment . }
+
+    # All member proteins of the same ComplexPortal complex
+    ?member a up:Protein ;
+            rdfs:seeAlso ?cp_db .
+    }
+    GROUP BY ?protein ?cp_db ?cp_comment
+    ORDER BY ?protein ?cp_db
+    LIMIT 500
+    ```
+
+    """
+    select_clause = dedent("""\
+        ?protein ?cp_db ?cp_comment
+        (GROUP_CONCAT(DISTINCT ?member; separator=",") AS ?complex_members)
+    """)
+    where_clause = dedent("""
+        # --- Complex Info ---
+        ?protein a up:Protein ;
+                rdfs:seeAlso ?cp_db .
+        ?cp_db up:database <http://purl.uniprot.org/database/ComplexPortal> .
+        OPTIONAL { ?cp_db rdfs:comment ?cp_comment . }
+        # All member proteins of the same ComplexPortal complex
+        ?member a up:Protein ;
+        rdfs:seeAlso ?cp_db .
+    """)
+    group_by = dedent("""
+       ?protein ?cp_db ?cp_comment
+    """)
+    return _build_sparql_generic_by_uniprot_accessions_query(
+        uniprot_accs, select_clause, where_clause, limit, groupby_clause=group_by
+    )
+
+
+@dataclass(frozen=True)
+class ComplexPortalEntry:
+    """A ComplexPortal entry.
+
+    Parameters:
+        query_protein: The UniProt accession used to find entry.
+        complex_id: The ComplexPortal identifier (for example "CPX-1234").
+        complex_url: The URL to the ComplexPortal entry.
+        complex_title: The title of the complex.
+        members: UniProt accessions which are members of the complex.
+    """
+
+    query_protein: str
+    complex_id: str
+    complex_url: str
+    complex_title: str
+    members: set[str]
+
+
+def _flatten_results_complex(raw_results) -> list[ComplexPortalEntry]:
+    results = []
+    for raw_result in raw_results:
+        query_protein = raw_result["protein"]["value"].split("/")[-1]
+        complex_id = raw_result["cp_db"]["value"].split("/")[-1]
+        complex_url = f"https://www.ebi.ac.uk/complexportal/complex/{complex_id}"
+        complex_title = raw_result.get("cp_comment", {}).get("value", "")
+        members = {m.split("/")[-1] for m in raw_result["complex_members"]["value"].split(",")}
+        results.append(
+            ComplexPortalEntry(
+                query_protein=query_protein,
+                complex_id=complex_id,
+                complex_url=complex_url,
+                complex_title=complex_title,
+                members=members,
+            )
+        )
+    return results
+
+
+def search4macromolecular_complexes(
+    uniprot_accs: Iterable[str], limit: int = 10_000, timeout: int = 1_800
+) -> list[ComplexPortalEntry]:
+    """Search for macromolecular complexes by UniProtKB accessions.
+
+    Queries for references to/from https://www.ebi.ac.uk/complexportal/ database in the Uniprot SPARQL endpoint.
+
+    Args:
+        uniprot_accs: UniProt accessions.
+        limit: Maximum number of results to return.
+        timeout: Timeout for the SPARQL query in seconds.
+
+    Returns:
+        List of ComplexPortalEntry objects.
+    """
+    sparql_query = _build_complex_sparql_query(uniprot_accs, limit)
+    logger.info("Executing SPARQL query for macromolecular complexes: %s", sparql_query)
+    raw_results = _execute_sparql_search(
+        sparql_query=sparql_query,
+        timeout=timeout,
+    )
+    limit_check("Search for complexes", limit, len(raw_results))
+    return _flatten_results_complex(raw_results)
+
+
+def search4interaction_partners(
+    uniprot_acc: str, excludes: set[str] | None = None, limit: int = 10_000, timeout: int = 1_800
+) -> dict[str, set[str]]:
+    """Search for interaction partners of a given UniProt accession using ComplexPortal database references.
+
+    Args:
+        uniprot_acc: UniProt accession to search interaction partners for.
+        excludes: Set of UniProt accessions to exclude from the results.
+            For example already known interaction partners.
+            If None then no complex members are excluded.
+        limit: Maximum number of results to return.
+        timeout: Timeout for the SPARQL query in seconds.
+
+    Returns:
+        Dictionary with UniProt accessions of interaction partners as keys and sets of ComplexPortal entry IDs
+        in which the interaction occurs as values.
+    """
+    ucomplexes = search4macromolecular_complexes([uniprot_acc], limit=limit, timeout=timeout)
+    hits: dict[str, set[str]] = {}
+    if excludes is None:
+        excludes = set()
+    for ucomplex in ucomplexes:
+        for member in ucomplex.members:
+            if member != uniprot_acc and member not in excludes:
+                if member not in hits:
+                    hits[member] = set()
+                hits[member].add(ucomplex.complex_id)
+    return hits

protein_quest/utils.py

@@ -1,22 +1,260 @@
 """Module for functions that are used in multiple places."""
 
+import argparse
 import asyncio
+import hashlib
 import logging
 import shutil
-from collections.abc import Coroutine, Iterable
+from collections.abc import Coroutine, Iterable, Sequence
 from contextlib import asynccontextmanager
+from functools import lru_cache
 from pathlib import Path
 from textwrap import dedent
-from typing import Any, Literal, get_args
+from typing import Any, Literal, Protocol, get_args, runtime_checkable
 
 import aiofiles
+import aiofiles.os
 import aiohttp
+import rich
+from aiohttp.streams import AsyncStreamIterator
 from aiohttp_retry import ExponentialRetry, RetryClient
+from platformdirs import user_cache_dir
+from rich_argparse import ArgumentDefaultsRichHelpFormatter
 from tqdm.asyncio import tqdm
 from yarl import URL
 
 logger = logging.getLogger(__name__)
 
+CopyMethod = Literal["copy", "symlink", "hardlink"]
+"""Methods for copying files."""
+copy_methods = set(get_args(CopyMethod))
+"""Set of valid copy methods."""
+
+
+@lru_cache
+def _cache_sub_dir(root_cache_dir: Path, filename: str, hash_length: int = 4) -> Path:
+    """Get the cache sub-directory for a given path.
+
+    To not have too many files in a single directory,
+    we create sub-directories based on the hash of the filename.
+
+    Args:
+        root_cache_dir: The root directory for the cache.
+        filename: The filename to be cached.
+        hash_length: The length of the hash to use for the sub-directory.
+
+    Returns:
+        The parent path to the cached file.
+    """
+    full_hash = hashlib.blake2b(filename.encode("utf-8")).hexdigest()
+    cache_sub_dir = full_hash[:hash_length]
+    cache_sub_dir_path = root_cache_dir / cache_sub_dir
+    cache_sub_dir_path.mkdir(parents=True, exist_ok=True)
+    return cache_sub_dir_path
+
+
+@runtime_checkable
+class Cacher(Protocol):
+    """Protocol for a cacher."""
+
+    def __contains__(self, item: str | Path) -> bool:
+        """Check if a file is in the cache.
+
+        Args:
+            item: The filename or Path to check.
+
+        Returns:
+            True if the file is in the cache, False otherwise.
+        """
+        ...
+
+    async def copy_from_cache(self, target: Path) -> Path | None:
+        """Copy a file from the cache to a target location if it exists in the cache.
+
+        Assumes:
+
+        - target does not exist.
+        - the parent directory of target exists.
+
+        Args:
+            target: The path to copy the file to.
+
+        Returns:
+            The path to the cached file if it was copied, None otherwise.
+        """
+        ...
+
+    async def write_iter(self, target: Path, content: AsyncStreamIterator[bytes]) -> Path:
+        """Write content to a file and cache it.
+
+        Args:
+            target: The path to write the content to.
+            content: An async iterator that yields bytes to write to the file.
+
+        Returns:
+            The path to the cached file.
+
+        Raises:
+            FileExistsError: If the target file already exists.
+        """
+        ...
+
+    async def write_bytes(self, target: Path, content: bytes) -> Path:
+        """Write bytes to a file and cache it.
+
+        Args:
+            target: The path to write the content to.
+            content: The bytes to write to the file.
+
+        Returns:
+            The path to the cached file.
+
+        Raises:
+            FileExistsError: If the target file already exists.
+        """
+        ...
+
+
+class PassthroughCacher(Cacher):
+    """A cacher that caches nothing.
+
+    On writes it just writes to the target path.
+    """
+
+    def __contains__(self, item: str | Path) -> bool:
+        # We don't have anything cached ever
+        return False
+
+    async def copy_from_cache(self, target: Path) -> Path | None:  # noqa: ARG002
+        # We don't have anything cached ever
+        return None
+
+    async def write_iter(self, target: Path, content: AsyncStreamIterator[bytes]) -> Path:
+        if target.exists():
+            raise FileExistsError(target)
+        target.write_bytes(b"".join([chunk async for chunk in content]))
+        return target
+
+    async def write_bytes(self, target: Path, content: bytes) -> Path:
+        if target.exists():
+            raise FileExistsError(target)
+        target.write_bytes(content)
+        return target
+
+
+def user_cache_root_dir() -> Path:
+    """Get the users root directory for caching files.
+
+    Returns:
+        The path to the user's cache directory for protein-quest.
+    """
+    return Path(user_cache_dir("protein-quest"))
+
+
+class DirectoryCacher(Cacher):
+    """Class to cache files in a directory.
+
+    Caching logic is based on the file name only.
+    If file name of paths are the same then the files are considered the same.
+
+    Attributes:
+        cache_dir: The directory to use for caching.
+        copy_method: The method to use for copying files.
+    """
+
+    def __init__(
+        self,
+        cache_dir: Path | None = None,
+        copy_method: CopyMethod = "hardlink",
+    ) -> None:
+        """Initialize the cacher.
+
+        If file name of paths are the same then the files are considered the same.
+
+        Args:
+            cache_dir: The directory to use for caching.
+                If None, a default cache directory (~/.cache/protein-quest) is used.
+            copy_method: The method to use for copying.
+        """
+        if cache_dir is None:
+            cache_dir = user_cache_root_dir()
+        self.cache_dir: Path = cache_dir
+        self.cache_dir.mkdir(parents=True, exist_ok=True)
+        if copy_method == "copy":
+            logger.warning(
+                "Using copy as copy_method to cache files is not recommended. "
+                "This will use more disk space and be slower than symlink or hardlink."
+            )
+        if copy_method not in copy_methods:
+            msg = f"Unknown copy method: {copy_method}. Must be one of {copy_methods}."
+            raise ValueError(msg)
+        self.copy_method: CopyMethod = copy_method
+
+    def __contains__(self, item: str | Path) -> bool:
+        cached_file = self._as_cached_path(item)
+        return cached_file.exists()
+
+    def _as_cached_path(self, item: str | Path) -> Path:
+        file_name = item.name if isinstance(item, Path) else item
+        cache_sub_dir = _cache_sub_dir(self.cache_dir, file_name)
+        return cache_sub_dir / file_name
+
+    async def copy_from_cache(self, target: Path) -> Path | None:
+        cached_file = self._as_cached_path(target.name)
+        exists = await aiofiles.os.path.exists(str(cached_file))
+        if exists:
+            await async_copyfile(cached_file, target, copy_method=self.copy_method)
+            return cached_file
+        return None
+
+    async def write_iter(self, target: Path, content: AsyncStreamIterator[bytes]) -> Path:
+        cached_file = self._as_cached_path(target.name)
+        # Write file to cache dir
+        async with aiofiles.open(cached_file, "xb") as f:
+            async for chunk in content:
+                await f.write(chunk)
+        # Copy to target location
+        await async_copyfile(cached_file, target, copy_method=self.copy_method)
+        return cached_file
+
+    async def write_bytes(self, target: Path, content: bytes) -> Path:
+        cached_file = self._as_cached_path(target.name)
+        # Write file to cache dir
+        async with aiofiles.open(cached_file, "xb") as f:
+            await f.write(content)
+        # Copy to target location
+        await async_copyfile(cached_file, target, copy_method=self.copy_method)
+        return cached_file
+
+    def populate_cache(self, source_dir: Path) -> dict[Path, Path]:
+        """Populate the cache from an existing directory.
+
+        This will copy all files from the source directory to the cache directory.
+        If a file with the same name already exists in the cache, it will be skipped.
+
+        Args:
+            source_dir: The directory to populate the cache from.
+
+        Returns:
+            A dictionary mapping source file paths to their cached paths.
+
+        Raises:
+            NotADirectoryError: If the source_dir is not a directory.
+        """
+        if not source_dir.is_dir():
+            raise NotADirectoryError(source_dir)
+        cached = {}
+        for file_path in source_dir.iterdir():
+            if not file_path.is_file():
+                continue
+            cached_path = self._as_cached_path(file_path.name)
+            if cached_path.exists():
+                logger.debug(f"File {file_path.name} already in cache. Skipping.")
+                continue
+            copyfile(file_path, cached_path, copy_method=self.copy_method)
+            cached[file_path] = cached_path
+        return cached
+
 
 async def retrieve_files(
     urls: Iterable[tuple[URL | str, str]],
@@ -25,6 +263,8 @@ async def retrieve_files(
     retries: int = 3,
     total_timeout: int = 300,
     desc: str = "Downloading files",
+    cacher: Cacher | None = None,
+    chunk_size: int = 524288,  # 512 KiB
 ) -> list[Path]:
     """Retrieve files from a list of URLs and save them to a directory.
 
@@ -35,6 +275,8 @@ async def retrieve_files(
         retries: The number of times to retry a failed download.
         total_timeout: The total timeout for a download in seconds.
         desc: Description for the progress bar.
+        cacher: An optional cacher to use for caching files.
+        chunk_size: The size of each chunk to read from the response.
 
     Returns:
         A list of paths to the downloaded files.
@@ -42,7 +284,17 @@ async def retrieve_files(
     save_dir.mkdir(parents=True, exist_ok=True)
     semaphore = asyncio.Semaphore(max_parallel_downloads)
     async with friendly_session(retries, total_timeout) as session:
-        tasks = [_retrieve_file(session, url, save_dir / filename, semaphore) for url, filename in urls]
+        tasks = [
+            _retrieve_file(
+                session=session,
+                url=url,
+                save_path=save_dir / filename,
+                semaphore=semaphore,
+                cacher=cacher,
+                chunk_size=chunk_size,
+            )
+            for url, filename in urls
+        ]
         files: list[Path] = await tqdm.gather(*tasks, desc=desc)
         return files
 
@@ -52,8 +304,8 @@ async def _retrieve_file(
     url: URL | str,
     save_path: Path,
     semaphore: asyncio.Semaphore,
-    ovewrite: bool = False,
-    chunk_size: int = 131072,  # 128 KiB
+    cacher: Cacher | None = None,
+    chunk_size: int = 524288,  # 512 KiB
 ) -> Path:
     """Retrieve a single file from a URL and save it to a specified path.
 
@@ -62,26 +314,28 @@ async def _retrieve_file(
         url: The URL to download the file from.
         save_path: The path where the file should be saved.
         semaphore: A semaphore to limit the number of concurrent downloads.
-        ovewrite: Whether to overwrite the file if it already exists.
+        cacher: An optional cacher to use for caching files.
         chunk_size: The size of each chunk to read from the response.
 
     Returns:
         The path to the saved file.
     """
     if save_path.exists():
-        if ovewrite:
-            save_path.unlink()
-        else:
-            logger.debug(f"File {save_path} already exists. Skipping download from {url}.")
-            return save_path
+        logger.debug(f"File {save_path} already exists. Skipping download from {url}.")
+        return save_path
+
+    if cacher is None:
+        cacher = PassthroughCacher()
+    if cached_file := await cacher.copy_from_cache(save_path):
+        logger.debug(f"File {save_path} was copied from cache {cached_file}. Skipping download from {url}.")
+        return save_path
+
     async with (
         semaphore,
-        aiofiles.open(save_path, "xb") as f,
         session.get(url) as resp,
     ):
         resp.raise_for_status()
-        async for chunk in resp.content.iter_chunked(chunk_size):
-            await f.write(chunk)
+        await cacher.write_iter(save_path, resp.content.iter_chunked(chunk_size))
     return save_path
 
 
@@ -141,27 +395,117 @@ def run_async[R](coroutine: Coroutine[Any, Any, R]) -> R:
         raise NestedAsyncIOLoopError from e
 
 
-CopyMethod = Literal["copy", "symlink"]
-copy_methods = set(get_args(CopyMethod))
-
-
 def copyfile(source: Path, target: Path, copy_method: CopyMethod = "copy"):
-    """Make target path be same file as source by either copying or symlinking.
+    """Make target path be same file as source by either copying or symlinking or hardlinking.
+
+    Note that the hardlink copy method only works within the same filesystem and is harder to track.
+    If you want to track cached files easily then use 'symlink'.
+    On Windows you need developer mode or admin privileges to create symlinks.
 
     Args:
-        source: The source file to copy or symlink.
+        source: The source file to copy or link.
         target: The target file to create.
         copy_method: The method to use for copying.
 
     Raises:
         FileNotFoundError: If the source file or parent of target does not exist.
-        ValueError: If the method is not "copy" or "symlink".
+        FileExistsError: If the target file already exists.
+        ValueError: If an unknown copy method is provided.
     """
     if copy_method == "copy":
         shutil.copyfile(source, target)
     elif copy_method == "symlink":
-        rel_source = source.relative_to(target.parent, walk_up=True)
+        rel_source = source.absolute().relative_to(target.parent.absolute(), walk_up=True)
         target.symlink_to(rel_source)
+    elif copy_method == "hardlink":
+        target.hardlink_to(source)
     else:
-        msg = f"Unknown method: {copy_method}"
+        msg = f"Unknown method: {copy_method}. Valid methods are: {copy_methods}"
         raise ValueError(msg)
+
+
+async def async_copyfile(
+    source: Path,
+    target: Path,
+    copy_method: CopyMethod = "copy",
+):
+    """Asynchronously make target path be same file as source by either copying or symlinking or hardlinking.
+
+    Note that the hardlink copy method only works within the same filesystem and is harder to track.
+    If you want to track cached files easily then use 'symlink'.
+    On Windows you need developer mode or admin privileges to create symlinks.
+
+    Args:
+        source: The source file to copy.
+        target: The target file to create.
+        copy_method: The method to use for copying.
+
+    Raises:
+        FileNotFoundError: If the source file or parent of target does not exist.
+        FileExistsError: If the target file already exists.
+        ValueError: If an unknown copy method is provided.
+    """
+    if copy_method == "copy":
+        # Could use loop of chunks with aiofiles,
+        # but shutil is ~1.9x faster on my machine
+        # due to fastcopy and sendfile optimizations in shutil.
+        await asyncio.to_thread(shutil.copyfile, source, target)
+    elif copy_method == "symlink":
+        rel_source = source.relative_to(target.parent, walk_up=True)
+        await aiofiles.os.symlink(str(rel_source), str(target))
+    elif copy_method == "hardlink":
+        await aiofiles.os.link(str(source), str(target))
+    else:
+        msg = f"Unknown method: {copy_method}. Valid methods are: {copy_methods}"
+        raise ValueError(msg)
+
+
+def populate_cache_command(raw_args: Sequence[str] | None = None):
+    """Command line interface to populate the cache from an existing directory.
+
+    Can be called from the command line as:
+
+    ```bash
+    python3 -m protein_quest.utils populate-cache /path/to/source/dir
+    ```
+
+    Args:
+        raw_args: The raw command line arguments to parse. If None, uses sys.argv.
+    """
+    root_parser = argparse.ArgumentParser(formatter_class=ArgumentDefaultsRichHelpFormatter)
+    subparsers = root_parser.add_subparsers(dest="command")
+
+    desc = "Populate the cache directory with files from the source directory."
+    populate_cache_parser = subparsers.add_parser(
+        "populate-cache",
+        help=desc,
+        description=desc,
+        formatter_class=ArgumentDefaultsRichHelpFormatter,
+    )
+    populate_cache_parser.add_argument("source_dir", type=Path)
+    populate_cache_parser.add_argument(
+        "--cache-dir",
+        type=Path,
+        default=user_cache_root_dir(),
+        help="Directory to use for caching. If not provided, a default cache directory is used.",
+    )
+    populate_cache_parser.add_argument(
+        "--copy-method",
+        type=str,
+        default="hardlink",
+        choices=copy_methods,
+        help="Method to use for copying files to cache.",
+    )
+
+    args = root_parser.parse_args(raw_args)
+    if args.command == "populate-cache":
+        source_dir = args.source_dir
+        cacher = DirectoryCacher(cache_dir=args.cache_dir, copy_method=args.copy_method)
+        cached_files = cacher.populate_cache(source_dir)
+        rich.print(f"Cached {len(cached_files)} files from {source_dir} to {cacher.cache_dir}")
+        for src, cached in cached_files.items():
+            rich.print(f"- {src.relative_to(source_dir)} -> {cached.relative_to(cacher.cache_dir)}")
+
+
+if __name__ == "__main__":
+    populate_cache_command()