refactor: lift files away from directory.py

Author: tristan-f-r

cache/Snakefile

@@ -1,34 +1,33 @@
-from cache import link
+from cache import FetchConfig, link
+from cache.directory import CacheItem # for exposing to Snakefiles that import this Snakefile.
 from cache.util import uncompress
 import urllib.parse
-from dataclasses import dataclass
 from typing import Union
 from pathlib import Path
 
-@dataclass
-class FetchConfig:
-    directive: list[str]
-    uncompress: bool = False
+def stringify_directive(directive: Union[CacheItem, FetchConfig]) -> str:
+    return urllib.parse.quote_plus(directive.name if isinstance(directive, CacheItem) else '/'.join(directive))
 
 def produce_fetch_rules(input_dict: dict[str, Union[FetchConfig, list[str]]]):
     """
     Produces fetch rules based on a dictionary mapping
     output files to their directory.py-based directive.
     """
-    # Map inputs to be wrapped with FetchConfig if list[str]
-    input_dict = {k: FetchConfig(v) if isinstance(v, list) else v for k, v in input_dict.items()}
+    # Map inputs to be wrapped with FetchConfig if list[str] or CacheItem
+    input_dict = {k: FetchConfig(v) if isinstance(v, tuple) or isinstance(v, CacheItem) else v for k, v in input_dict.items()}
 
-    directives = [urllib.parse.quote_plus("/".join(directive.directive)) for directive in input_dict.values()]
+    directives = list(input_dict.values())
     assert len(directives) == len(set(directives)), "Directives aren't unique!"
 
     for output_file, config in input_dict.items():
         # Since placeholders are evaluated when the job is actually ran,
         # we pass data using params and output.
         rule:
-            name: f"fetch_{urllib.parse.quote_plus('/'.join(config.directive))}_to_{urllib.parse.quote_plus(output_file)}"
+            name:
+                f"fetch_{stringify_directive(config.directive)}_to_{urllib.parse.quote_plus(output_file)}"
             output: file=output_file
             params:
                 config=config
             run:
                 Path(output.file).parent.mkdir(exist_ok=True)
-                link(Path(output.file), params.config.directive, uncompress=params.config.uncompress)
+                link(Path(output.file), params.config)

cache/__init__.py

@@ -2,34 +2,41 @@
 This is how spras-benchmarking handles artifact caching. `cache` should be used specifically inside `Snakefile`
 """
 
+from dataclasses import dataclass
+from typing import Union
 from cache.util import uncompress as uncompress_file
-from cache.directory import get_cache_item
+from cache.directory import CacheItem, get_cache_item
 from pathlib import Path
 import os
 from urllib.parse import quote_plus
 import pickle
 
-__all__ = ["link"]
+__all__ = ["FetchConfig", "link"]
 
 dir_path = Path(os.path.dirname(os.path.realpath(__file__)))
 artifacts_dir = dir_path / "artifacts"
 
+@dataclass(frozen=True)
+class FetchConfig:
+    directive: Union[CacheItem, tuple[str, ...]]
+    uncompress: bool = False
 
-def get_artifact_name(directive: list[str]) -> str:
+def get_artifact_name(directive: tuple[str, ...]) -> str:
     return quote_plus("/".join(directive))
 
+def add_suffix(path: Path, suffix: str):
+    return path.with_suffix(path.suffix + suffix)
 
-def has_expired(directive: list[str]) -> bool:
+def has_expired(
+        cache_item: CacheItem,
+        output: Path
+) -> bool:
     """
     Check if the artifact metadata associated with a directive has expired.
     Avoids re-downloading the artifact if nothing has changed.
     """
-    artifact_name = get_artifact_name(directive)
-    cache_item = get_cache_item(directive)
 
-    metadata_dir = artifacts_dir / "metadata"
-    metadata_dir.mkdir(exist_ok=True)
-    metadata_file = (artifacts_dir / "metadata" / artifact_name).with_suffix((artifacts_dir / artifact_name).suffix + ".metadata")
+    metadata_file = add_suffix(output, ".metadata")
 
     # metadata never existed: we need to retrieve the new file
     if not metadata_file.exists():
@@ -50,36 +57,71 @@ def has_expired(directive: list[str]) -> bool:
     # metadata hasn't changed and already existed: this hasn't expired
     return False
 
+def link_with_cache_item(
+    output: Path,
+    cache_item: CacheItem,
+    uncompress: bool = False
+):
+    """
+    Intermediary function for `link`.
+    This does almost all of what `link` is characterized to do in its documentation,
+    except for doing symlinking.
+    """
+    # If `uncompress` is `True`, we make
+    # `output` our 'compressed output.'
+    uncompressed_output = output
+    if uncompress:
+        output = add_suffix(output, ".compresseded")
+
+    # Re-download if the file doesn't exist or the directive has expired.
+    # Note that we check for expiration first to trigger metadata creation.
+    if has_expired(cache_item, output) or not output.exists():
+        output.unlink(missing_ok=True)
+        cache_item.download(output)
+
+    if uncompress:
+        uncompressed_artifact_path = add_suffix(output, ".uncompressed")
+        uncompressed_artifact_path.unlink(missing_ok=True)
+        uncompress_file(output, uncompressed_output)
 
-def link(output: str, directive: list[str], uncompress=False):
+def link(
+        output: str,
+        config: FetchConfig
+):
     """
     Links output files from cache.directory directives.
     For example,
 
     ```py
-    link("output/ensg-ensp.tsv", ["BioMart", "ensg-ensp.tsv"])
+    link("output/ensg-ensp.tsv", FetchConfig(["BioMart", "ensg-ensp.tsv"]))
     ```
 
-    would download and check BioMart's cache for ENSG-ENSP mapping, then symlink the cached output
-    (lying somewhere in the cache folder) with the desired `output`.
-    """
-
-    artifacts_dir.mkdir(exist_ok=True)
-
-    artifact_name = get_artifact_name(directive)
-
-    Path(output).unlink(missing_ok=True)
+    would download and check BioMart's cache for ENSG-ENSP mapping, then:
+    - If `config.directive` is a `CacheItem`, we write the file directly to `output`.
+    - Otherwise, we symlink the cached output (lying somewhere in the cache folder) with the desired `output`
+    to avoid file duplication.
 
-    # Re-download if the file doesn't exist or the directive has expired.
-    cache_item = get_cache_item(directive)
-    if not (artifacts_dir / artifact_name).exists() or has_expired(directive):
-        (artifacts_dir / artifact_name).unlink(missing_ok=True)
-        cache_item.download(artifacts_dir / artifact_name)
+    This function wraps around link_with_cache_item and handles symlinking
+    depending on the type of config.directive.
+    TODO: most likely a nicer way to design this.
+    """
 
-    if uncompress:
-        uncompressed_artifact_path = Path(str(artifacts_dir / artifact_name) + ".uncompressed")
-        uncompressed_artifact_path.unlink(missing_ok=True)
-        uncompress_file(artifacts_dir / artifact_name, uncompressed_artifact_path)
-        Path(output).symlink_to(uncompressed_artifact_path)
+    if isinstance(config.directive, CacheItem):
+        link_with_cache_item(
+            Path(output),
+            config.directive,
+            config.uncompress
+        )
     else:
-        Path(output).symlink_to(artifacts_dir / artifact_name)
+        artifacts_dir.mkdir(exist_ok=True)
+        artifact_name = get_artifact_name(config.directive)
+        artifact_output = artifacts_dir / artifact_name
+
+        link_with_cache_item(
+            artifact_output,
+            get_cache_item(config.directive),
+            config.uncompress
+        )
+
+        Path(output).symlink_to(artifact_output)
+