codekingpro/portable-devtools

1import os2import re3import time4from collections.abc import Iterator, Sequence5from pathlib import Path6 7from langchain_core.stores import ByteStore8 9from langchain_classic.storage.exceptions import InvalidKeyException10 11 12class LocalFileStore(ByteStore):13    """`BaseStore` interface that works on the local file system.14 15    Examples:16        Create a `LocalFileStore` instance and perform operations on it:17 18        ```python19        from langchain_classic.storage import LocalFileStore20 21        # Instantiate the LocalFileStore with the root path22        file_store = LocalFileStore("/path/to/root")23 24        # Set values for keys25        file_store.mset([("key1", b"value1"), ("key2", b"value2")])26 27        # Get values for keys28        values = file_store.mget(["key1", "key2"])  # Returns [b"value1", b"value2"]29 30        # Delete keys31        file_store.mdelete(["key1"])32 33        # Iterate over keys34        for key in file_store.yield_keys():35            print(key)  # noqa: T20136        ```37    """38 39    def __init__(40        self,41        root_path: str | Path,42        *,43        chmod_file: int | None = None,44        chmod_dir: int | None = None,45        update_atime: bool = False,46    ) -> None:47        """Implement the `BaseStore` interface for the local file system.48 49        Args:50            root_path: The root path of the file store. All keys are interpreted as51                paths relative to this root.52            chmod_file: Sets permissions for newly created files, overriding the53                current `umask` if needed.54            chmod_dir: Sets permissions for newly created dirs, overriding the55                current `umask` if needed.56            update_atime: Updates the filesystem access time (but not the modified57                time) when a file is read. This allows MRU/LRU cache policies to be58                implemented for filesystems where access time updates are disabled.59        """60        self.root_path = Path(root_path).absolute()61        self.chmod_file = chmod_file62        self.chmod_dir = chmod_dir63        self.update_atime = update_atime64 65    def _get_full_path(self, key: str) -> Path:66        """Get the full path for a given key relative to the root path.67 68        Args:69            key: The key relative to the root path.70 71        Returns:72            The full path for the given key.73        """74        if not re.match(r"^[a-zA-Z0-9_.\-/]+$", key):75            msg = f"Invalid characters in key: {key}"76            raise InvalidKeyException(msg)77        full_path = (self.root_path / key).resolve()78        root_path = self.root_path.resolve()79        common_path = os.path.commonpath([root_path, full_path])80        if common_path != str(root_path):81            msg = (82                f"Invalid key: {key}. Key should be relative to the full path. "83                f"{root_path} vs. {common_path} and full path of {full_path}"84            )85            raise InvalidKeyException(msg)86 87        return full_path88 89    def _mkdir_for_store(self, dir_path: Path) -> None:90        """Makes a store directory path (including parents) with specified permissions.91 92        This is needed because `Path.mkdir()` is restricted by the current `umask`,93        whereas the explicit `os.chmod()` used here is not.94 95        Args:96            dir_path: The store directory to make.97        """98        if not dir_path.exists():99            self._mkdir_for_store(dir_path.parent)100            dir_path.mkdir(exist_ok=True)101            if self.chmod_dir is not None:102                dir_path.chmod(self.chmod_dir)103 104    def mget(self, keys: Sequence[str]) -> list[bytes | None]:105        """Get the values associated with the given keys.106 107        Args:108            keys: A sequence of keys.109 110        Returns:111            A sequence of optional values associated with the keys.112            If a key is not found, the corresponding value will be `None`.113        """114        values: list[bytes | None] = []115        for key in keys:116            full_path = self._get_full_path(key)117            if full_path.exists():118                value = full_path.read_bytes()119                values.append(value)120                if self.update_atime:121                    # update access time only; preserve modified time122                    os.utime(full_path, (time.time(), full_path.stat().st_mtime))123            else:124                values.append(None)125        return values126 127    def mset(self, key_value_pairs: Sequence[tuple[str, bytes]]) -> None:128        """Set the values for the given keys.129 130        Args:131            key_value_pairs: A sequence of key-value pairs.132        """133        for key, value in key_value_pairs:134            full_path = self._get_full_path(key)135            self._mkdir_for_store(full_path.parent)136            full_path.write_bytes(value)137            if self.chmod_file is not None:138                full_path.chmod(self.chmod_file)139 140    def mdelete(self, keys: Sequence[str]) -> None:141        """Delete the given keys and their associated values.142 143        Args:144            keys: A sequence of keys to delete.145        """146        for key in keys:147            full_path = self._get_full_path(key)148            if full_path.exists():149                full_path.unlink()150 151    def yield_keys(self, *, prefix: str | None = None) -> Iterator[str]:152        """Get an iterator over keys that match the given prefix.153 154        Args:155            prefix: The prefix to match.156 157        Yields:158            Keys that match the given prefix.159        """160        prefix_path = self._get_full_path(prefix) if prefix else self.root_path161        for file in prefix_path.rglob("*"):162            if file.is_file():163                relative_path = file.relative_to(self.root_path)164                yield str(relative_path)165