src/black/cache.py

   1 """Caching of formatted files with feature-based invalidation."""
   2
   3 import os
   4 import pickle
   5 from pathlib import Path
   6 import tempfile
   7 from typing import Dict, Iterable, Set, Tuple
   8
   9 from platformdirs import user_cache_dir
  10
  11 from black.mode import Mode
  12
  13 from _black_version import version as __version__
  14
  15
  16 # types
  17 Timestamp = float
  18 FileSize = int
  19 CacheInfo = Tuple[Timestamp, FileSize]
  20 Cache = Dict[str, CacheInfo]
  21
  22
  23 def get_cache_dir() -> Path:
  24     """Get the cache directory used by black.
  25
  26     Users can customize this directory on all systems using `BLACK_CACHE_DIR`
  27     environment variable. By default, the cache directory is the user cache directory
  28     under the black application.
  29
  30     This result is immediately set to a constant `black.cache.CACHE_DIR` as to avoid
  31     repeated calls.
  32     """
  33     # NOTE: Function mostly exists as a clean way to test getting the cache directory.
  34     default_cache_dir = user_cache_dir("black", version=__version__)
  35     cache_dir = Path(os.environ.get("BLACK_CACHE_DIR", default_cache_dir))
  36     return cache_dir
  37
  38
  39 CACHE_DIR = get_cache_dir()
  40
  41
  42 def read_cache(mode: Mode) -> Cache:
  43     """Read the cache if it exists and is well formed.
  44
  45     If it is not well formed, the call to write_cache later should resolve the issue.
  46     """
  47     cache_file = get_cache_file(mode)
  48     if not cache_file.exists():
  49         return {}
  50
  51     with cache_file.open("rb") as fobj:
  52         try:
  53             cache: Cache = pickle.load(fobj)
  54         except (pickle.UnpicklingError, ValueError, IndexError):
  55             return {}
  56
  57     return cache
  58
  59
  60 def get_cache_file(mode: Mode) -> Path:
  61     return CACHE_DIR / f"cache.{mode.get_cache_key()}.pickle"
  62
  63
  64 def get_cache_info(path: Path) -> CacheInfo:
  65     """Return the information used to check if a file is already formatted or not."""
  66     stat = path.stat()
  67     return stat.st_mtime, stat.st_size
  68
  69
  70 def filter_cached(cache: Cache, sources: Iterable[Path]) -> Tuple[Set[Path], Set[Path]]:
  71     """Split an iterable of paths in `sources` into two sets.
  72
  73     The first contains paths of files that modified on disk or are not in the
  74     cache. The other contains paths to non-modified files.
  75     """
  76     todo, done = set(), set()
  77     for src in sources:
  78         res_src = src.resolve()
  79         if cache.get(str(res_src)) != get_cache_info(res_src):
  80             todo.add(src)
  81         else:
  82             done.add(src)
  83     return todo, done
  84
  85
  86 def write_cache(cache: Cache, sources: Iterable[Path], mode: Mode) -> None:
  87     """Update the cache file."""
  88     cache_file = get_cache_file(mode)
  89     try:
  90         CACHE_DIR.mkdir(parents=True, exist_ok=True)
  91         new_cache = {
  92             **cache,
  93             **{str(src.resolve()): get_cache_info(src) for src in sources},
  94         }
  95         with tempfile.NamedTemporaryFile(dir=str(cache_file.parent), delete=False) as f:
  96             pickle.dump(new_cache, f, protocol=4)
  97         os.replace(f.name, cache_file)
  98     except OSError:
  99         pass