CoolFace
Apppublic

Aluode/PerceptionLabPortable

sourceHugging Faceupdated 9mo agoView on Hugging Face
0likes
memory.py312 linesDownload Raw Back to implementations
1from __future__ import annotations2 3import logging4from datetime import datetime, timezone5from errno import ENOTEMPTY6from io import BytesIO7from pathlib import PurePath, PureWindowsPath8from typing import Any, ClassVar9 10from fsspec import AbstractFileSystem11from fsspec.implementations.local import LocalFileSystem12from fsspec.utils import stringify_path13 14logger = logging.getLogger("fsspec.memoryfs")15 16 17class MemoryFileSystem(AbstractFileSystem):18    """A filesystem based on a dict of BytesIO objects19 20    This is a global filesystem so instances of this class all point to the same21    in memory filesystem.22    """23 24    store: ClassVar[dict[str, Any]] = {}  # global, do not overwrite!25    pseudo_dirs = [""]  # global, do not overwrite!26    protocol = "memory"27    root_marker = "/"28 29    @classmethod30    def _strip_protocol(cls, path):31        if isinstance(path, PurePath):32            if isinstance(path, PureWindowsPath):33                return LocalFileSystem._strip_protocol(path)34            else:35                path = stringify_path(path)36 37        path = path.removeprefix("memory://")38        if "::" in path or "://" in path:39            return path.rstrip("/")40        path = path.lstrip("/").rstrip("/")41        return "/" + path if path else ""42 43    def ls(self, path, detail=True, **kwargs):44        path = self._strip_protocol(path)45        if path in self.store:46            # there is a key with this exact name47            if not detail:48                return [path]49            return [50                {51                    "name": path,52                    "size": self.store[path].size,53                    "type": "file",54                    "created": self.store[path].created.timestamp(),55                }56            ]57        paths = set()58        starter = path + "/"59        out = []60        for p2 in tuple(self.store):61            if p2.startswith(starter):62                if "/" not in p2[len(starter) :]:63                    # exact child64                    out.append(65                        {66                            "name": p2,67                            "size": self.store[p2].size,68                            "type": "file",69                            "created": self.store[p2].created.timestamp(),70                        }71                    )72                elif len(p2) > len(starter):73                    # implied child directory74                    ppath = starter + p2[len(starter) :].split("/", 1)[0]75                    if ppath not in paths:76                        out = out or []77                        out.append(78                            {79                                "name": ppath,80                                "size": 0,81                                "type": "directory",82                            }83                        )84                        paths.add(ppath)85        for p2 in self.pseudo_dirs:86            if p2.startswith(starter):87                if "/" not in p2[len(starter) :]:88                    # exact child pdir89                    if p2 not in paths:90                        out.append({"name": p2, "size": 0, "type": "directory"})91                        paths.add(p2)92                else:93                    # directory implied by deeper pdir94                    ppath = starter + p2[len(starter) :].split("/", 1)[0]95                    if ppath not in paths:96                        out.append({"name": ppath, "size": 0, "type": "directory"})97                        paths.add(ppath)98        if not out:99            if path in self.pseudo_dirs:100                # empty dir101                return []102            raise FileNotFoundError(path)103        if detail:104            return out105        return sorted([f["name"] for f in out])106 107    def mkdir(self, path, create_parents=True, **kwargs):108        path = self._strip_protocol(path)109        if path in self.store or path in self.pseudo_dirs:110            raise FileExistsError(path)111        if self._parent(path).strip("/") and self.isfile(self._parent(path)):112            raise NotADirectoryError(self._parent(path))113        if create_parents and self._parent(path).strip("/"):114            try:115                self.mkdir(self._parent(path), create_parents, **kwargs)116            except FileExistsError:117                pass118        if path and path not in self.pseudo_dirs:119            self.pseudo_dirs.append(path)120 121    def makedirs(self, path, exist_ok=False):122        try:123            self.mkdir(path, create_parents=True)124        except FileExistsError:125            if not exist_ok:126                raise127 128    def pipe_file(self, path, value, mode="overwrite", **kwargs):129        """Set the bytes of given file130 131        Avoids copies of the data if possible132        """133        mode = "xb" if mode == "create" else "wb"134        self.open(path, mode=mode, data=value)135 136    def rmdir(self, path):137        path = self._strip_protocol(path)138        if path == "":139            # silently avoid deleting FS root140            return141        if path in self.pseudo_dirs:142            if not self.ls(path):143                self.pseudo_dirs.remove(path)144            else:145                raise OSError(ENOTEMPTY, "Directory not empty", path)146        else:147            raise FileNotFoundError(path)148 149    def info(self, path, **kwargs):150        logger.debug("info: %s", path)151        path = self._strip_protocol(path)152        if path in self.pseudo_dirs or any(153            p.startswith(path + "/") for p in list(self.store) + self.pseudo_dirs154        ):155            return {156                "name": path,157                "size": 0,158                "type": "directory",159            }160        elif path in self.store:161            filelike = self.store[path]162            return {163                "name": path,164                "size": filelike.size,165                "type": "file",166                "created": getattr(filelike, "created", None),167            }168        else:169            raise FileNotFoundError(path)170 171    def _open(172        self,173        path,174        mode="rb",175        block_size=None,176        autocommit=True,177        cache_options=None,178        **kwargs,179    ):180        path = self._strip_protocol(path)181        if "x" in mode and self.exists(path):182            raise FileExistsError183        if path in self.pseudo_dirs:184            raise IsADirectoryError(path)185        parent = path186        while len(parent) > 1:187            parent = self._parent(parent)188            if self.isfile(parent):189                raise FileExistsError(parent)190        if mode in ["rb", "ab", "r+b", "a+b"]:191            if path in self.store:192                f = self.store[path]193                if "a" in mode:194                    # position at the end of file195                    f.seek(0, 2)196                else:197                    # position at the beginning of file198                    f.seek(0)199                return f200            else:201                raise FileNotFoundError(path)202        elif mode in {"wb", "w+b", "xb", "x+b"}:203            if "x" in mode and self.exists(path):204                raise FileExistsError205            m = MemoryFile(self, path, kwargs.get("data"))206            if not self._intrans:207                m.commit()208            return m209        else:210            name = self.__class__.__name__211            raise ValueError(f"unsupported file mode for {name}: {mode!r}")212 213    def cp_file(self, path1, path2, **kwargs):214        path1 = self._strip_protocol(path1)215        path2 = self._strip_protocol(path2)216        if self.isfile(path1):217            self.store[path2] = MemoryFile(218                self, path2, self.store[path1].getvalue()219            )  # implicit copy220        elif self.isdir(path1):221            if path2 not in self.pseudo_dirs:222                self.pseudo_dirs.append(path2)223        else:224            raise FileNotFoundError(path1)225 226    def cat_file(self, path, start=None, end=None, **kwargs):227        logger.debug("cat: %s", path)228        path = self._strip_protocol(path)229        try:230            return bytes(self.store[path].getbuffer()[start:end])231        except KeyError as e:232            raise FileNotFoundError(path) from e233 234    def _rm(self, path):235        path = self._strip_protocol(path)236        try:237            del self.store[path]238        except KeyError as e:239            raise FileNotFoundError(path) from e240 241    def modified(self, path):242        path = self._strip_protocol(path)243        try:244            return self.store[path].modified245        except KeyError as e:246            raise FileNotFoundError(path) from e247 248    def created(self, path):249        path = self._strip_protocol(path)250        try:251            return self.store[path].created252        except KeyError as e:253            raise FileNotFoundError(path) from e254 255    def isfile(self, path):256        path = self._strip_protocol(path)257        return path in self.store258 259    def rm(self, path, recursive=False, maxdepth=None):260        if isinstance(path, str):261            path = self._strip_protocol(path)262        else:263            path = [self._strip_protocol(p) for p in path]264        paths = self.expand_path(path, recursive=recursive, maxdepth=maxdepth)265        for p in reversed(paths):266            if self.isfile(p):267                self.rm_file(p)268            # If the expanded path doesn't exist, it is only because the expanded269            # path was a directory that does not exist in self.pseudo_dirs. This270            # is possible if you directly create files without making the271            # directories first.272            elif not self.exists(p):273                continue274            else:275                self.rmdir(p)276 277 278class MemoryFile(BytesIO):279    """A BytesIO which can't close and works as a context manager280 281    Can initialise with data. Each path should only be active once at any moment.282 283    No need to provide fs, path if auto-committing (default)284    """285 286    def __init__(self, fs=None, path=None, data=None):287        logger.debug("open file %s", path)288        self.fs = fs289        self.path = path290        self.created = datetime.now(tz=timezone.utc)291        self.modified = datetime.now(tz=timezone.utc)292        if data:293            super().__init__(data)294            self.seek(0)295 296    @property297    def size(self):298        return self.getbuffer().nbytes299 300    def __enter__(self):301        return self302 303    def close(self):304        pass305 306    def discard(self):307        pass308 309    def commit(self):310        self.fs.store[self.path] = self311        self.modified = datetime.now(tz=timezone.utc)312 
Aluode/PerceptionLabPortable · CoolFace