codekingpro/portable-devtools

1from __future__ import annotations2 3import os4import pickle5import time6from typing import TYPE_CHECKING7 8from fsspec.utils import atomic_write9 10try:11    import ujson as json12except ImportError:13    if not TYPE_CHECKING:14        import json15 16if TYPE_CHECKING:17    from collections.abc import Iterator18    from typing import Any, Literal, TypeAlias19 20    from .cached import CachingFileSystem21 22    Detail: TypeAlias = dict[str, Any]23 24 25class CacheMetadata:26    """Cache metadata.27 28    All reading and writing of cache metadata is performed by this class,29    accessing the cached files and blocks is not.30 31    Metadata is stored in a single file per storage directory in JSON format.32    For backward compatibility, also reads metadata stored in pickle format33    which is converted to JSON when next saved.34    """35 36    def __init__(self, storage: list[str]):37        """38 39        Parameters40        ----------41        storage: list[str]42            Directories containing cached files, must be at least one. Metadata43            is stored in the last of these directories by convention.44        """45        if not storage:46            raise ValueError("CacheMetadata expects at least one storage location")47 48        self._storage = storage49        self.cached_files: list[Detail] = [{}]50 51        # Private attribute to force saving of metadata in pickle format rather than52        # JSON for use in tests to confirm can read both pickle and JSON formats.53        self._force_save_pickle = False54 55    def _load(self, fn: str) -> Detail:56        """Low-level function to load metadata from specific file"""57        try:58            with open(fn, "r") as f:59                loaded = json.load(f)60        except ValueError:61            with open(fn, "rb") as f:62                loaded = pickle.load(f)63        for c in loaded.values():64            if isinstance(c.get("blocks"), list):65                c["blocks"] = set(c["blocks"])66        return loaded67 68    def _save(self, metadata_to_save: Detail, fn: str) -> None:69        """Low-level function to save metadata to specific file"""70        if self._force_save_pickle:71            with atomic_write(fn) as f:72                pickle.dump(metadata_to_save, f)73        else:74            with atomic_write(fn, mode="w") as f:75                json.dump(metadata_to_save, f)76 77    def _scan_locations(78        self, writable_only: bool = False79    ) -> Iterator[tuple[str, str, bool]]:80        """Yield locations (filenames) where metadata is stored, and whether81        writable or not.82 83        Parameters84        ----------85        writable: bool86            Set to True to only yield writable locations.87 88        Returns89        -------90        Yields (str, str, bool)91        """92        n = len(self._storage)93        for i, storage in enumerate(self._storage):94            writable = i == n - 195            if writable_only and not writable:96                continue97            yield os.path.join(storage, "cache"), storage, writable98 99    def check_file(100        self, path: str, cfs: CachingFileSystem | None101    ) -> Literal[False] | tuple[Detail, str]:102        """If path is in cache return its details, otherwise return ``False``.103 104        If the optional CachingFileSystem is specified then it is used to105        perform extra checks to reject possible matches, such as if they are106        too old.107        """108        for (fn, base, _), cache in zip(self._scan_locations(), self.cached_files):109            if path not in cache:110                continue111            detail = cache[path].copy()112 113            if cfs is not None:114                if cfs.check_files and detail["uid"] != cfs.fs.ukey(path):115                    # Wrong file as determined by hash of file properties116                    continue117                if cfs.expiry and time.time() - detail["time"] > cfs.expiry:118                    # Cached file has expired119                    continue120 121            fn = os.path.join(base, detail["fn"])122            if os.path.exists(fn):123                return detail, fn124        return False125 126    def clear_expired(self, expiry_time: int) -> tuple[list[str], bool]:127        """Remove expired metadata from the cache.128 129        Returns names of files corresponding to expired metadata and a boolean130        flag indicating whether the writable cache is empty. Caller is131        responsible for deleting the expired files.132        """133        expired_files = []134        for path, detail in self.cached_files[-1].copy().items():135            if time.time() - detail["time"] > expiry_time:136                fn = detail.get("fn", "")137                if not fn:138                    raise RuntimeError(139                        f"Cache metadata does not contain 'fn' for {path}"140                    )141                fn = os.path.join(self._storage[-1], fn)142                expired_files.append(fn)143                self.cached_files[-1].pop(path)144 145        if self.cached_files[-1]:146            cache_path = os.path.join(self._storage[-1], "cache")147            self._save(self.cached_files[-1], cache_path)148 149        writable_cache_empty = not self.cached_files[-1]150        return expired_files, writable_cache_empty151 152    def load(self) -> None:153        """Load all metadata from disk and store in ``self.cached_files``"""154        cached_files = []155        for fn, _, _ in self._scan_locations():156            if os.path.exists(fn):157                # TODO: consolidate blocks here158                cached_files.append(self._load(fn))159            else:160                cached_files.append({})161        self.cached_files = cached_files or [{}]162 163    def on_close_cached_file(self, f: Any, path: str) -> None:164        """Perform side-effect actions on closing a cached file.165 166        The actual closing of the file is the responsibility of the caller.167        """168        # File must be writable, so in self.cached_files[-1]169        c = self.cached_files[-1][path]170        if c["blocks"] is not True and len(c["blocks"]) * f.blocksize >= f.size:171            c["blocks"] = True172 173    def pop_file(self, path: str) -> str | None:174        """Remove metadata of cached file.175 176        If path is in the cache, return the filename of the cached file,177        otherwise return ``None``.  Caller is responsible for deleting the178        cached file.179        """180        details = self.check_file(path, None)181        if not details:182            return None183        _, fn = details184        if fn.startswith(self._storage[-1]):185            self.cached_files[-1].pop(path)186            self.save()187        else:188            raise PermissionError(189                "Can only delete cached file in last, writable cache location"190            )191        return fn192 193    def save(self) -> None:194        """Save metadata to disk"""195        for (fn, _, writable), cache in zip(self._scan_locations(), self.cached_files):196            if not writable:197                continue198 199            if os.path.exists(fn):200                cached_files = self._load(fn)201                for k, c in cached_files.items():202                    if k in cache:203                        if c["blocks"] is True or cache[k]["blocks"] is True:204                            c["blocks"] = True205                        else:206                            # self.cached_files[*][*]["blocks"] must continue to207                            # point to the same set object so that updates208                            # performed by MMapCache are propagated back to209                            # self.cached_files.210                            blocks = cache[k]["blocks"]211                            blocks.update(c["blocks"])212                            c["blocks"] = blocks213                        c["time"] = max(c["time"], cache[k]["time"])214                        c["uid"] = cache[k]["uid"]215 216                # Files can be added to cache after it was written once217                for k, c in cache.items():218                    if k not in cached_files:219                        cached_files[k] = c220            else:221                cached_files = cache222            cache = {k: v.copy() for k, v in cached_files.items()}223            for c in cache.values():224                if isinstance(c["blocks"], set):225                    c["blocks"] = list(c["blocks"])226            self._save(cache, fn)227            self.cached_files[-1] = cached_files228 229    def update_file(self, path: str, detail: Detail) -> None:230        """Update metadata for specific file in memory, do not save"""231        self.cached_files[-1][path] = detail232