codekingpro/portable-devtools

1import os2import zipfile3 4import fsspec5from fsspec.archive import AbstractArchiveFileSystem6 7 8class ZipFileSystem(AbstractArchiveFileSystem):9    """Read/Write contents of ZIP archive as a file-system10 11    Keeps file object open while instance lives.12 13    This class is pickleable, but not necessarily thread-safe14    """15 16    root_marker = ""17    protocol = "zip"18    cachable = False19 20    def __init__(21        self,22        fo="",23        mode="r",24        target_protocol=None,25        target_options=None,26        compression=zipfile.ZIP_STORED,27        allowZip64=True,28        compresslevel=None,29        **kwargs,30    ):31        """32        Parameters33        ----------34        fo: str or file-like35            Contains ZIP, and must exist. If a str, will fetch file using36            :meth:`~fsspec.open_files`, which must return one file exactly.37        mode: str38            Accept: "r", "w", "a"39        target_protocol: str (optional)40            If ``fo`` is a string, this value can be used to override the41            FS protocol inferred from a URL42        target_options: dict (optional)43            Kwargs passed when instantiating the target FS, if ``fo`` is44            a string.45        compression, allowZip64, compresslevel: passed to ZipFile46            Only relevant when creating a ZIP47        """48        super().__init__(self, **kwargs)49        if mode not in set("rwa"):50            raise ValueError(f"mode '{mode}' no understood")51        self.mode = mode52        if isinstance(fo, (str, os.PathLike)):53            if mode == "a":54                m = "r+b"55            else:56                m = mode + "b"57            fo = fsspec.open(58                fo, mode=m, protocol=target_protocol, **(target_options or {})59            )60        self.force_zip_64 = allowZip6461        self.of = fo62        self.fo = fo.__enter__()  # the whole instance is a context63        self.zip = zipfile.ZipFile(64            self.fo,65            mode=mode,66            compression=compression,67            allowZip64=allowZip64,68            compresslevel=compresslevel,69        )70        self.dir_cache = None71 72    @classmethod73    def _strip_protocol(cls, path):74        # zip file paths are always relative to the archive root75        return super()._strip_protocol(path).lstrip("/")76 77    def __del__(self):78        if hasattr(self, "zip"):79            self.close()80            del self.zip81        if hasattr(self, "of") and hasattr(self.of, "__exit__"):82            self.of.__exit__(None, None, None)83 84    def close(self):85        """Commits any write changes to the file. Done on ``del`` too."""86        self.zip.close()87 88    def _get_dirs(self):89        if self.dir_cache is None or self.mode in set("wa"):90            # when writing, dir_cache is always in the ZipFile's attributes,91            # not read from the file.92            files = self.zip.infolist()93            self.dir_cache = {94                dirname.rstrip("/"): {95                    "name": dirname.rstrip("/"),96                    "size": 0,97                    "type": "directory",98                }99                for dirname in self._all_dirnames(self.zip.namelist())100            }101            for z in files:102                f = {s: getattr(z, s, None) for s in zipfile.ZipInfo.__slots__}103                f.update(104                    {105                        "name": z.filename.rstrip("/"),106                        "size": z.file_size,107                        "type": ("directory" if z.is_dir() else "file"),108                    }109                )110                self.dir_cache[f["name"]] = f111 112    def pipe_file(self, path, value, **kwargs):113        # override upstream, because we know the exact file size in this case114        self.zip.writestr(path, value, **kwargs)115 116    def _open(117        self,118        path,119        mode="rb",120        block_size=None,121        autocommit=True,122        cache_options=None,123        **kwargs,124    ):125        path = self._strip_protocol(path)126        if "r" in mode and self.mode in set("wa"):127            if self.exists(path):128                raise OSError("ZipFS can only be open for reading or writing, not both")129            raise FileNotFoundError(path)130        if "r" in self.mode and "w" in mode:131            raise OSError("ZipFS can only be open for reading or writing, not both")132        out = self.zip.open(path, mode.strip("b"), force_zip64=self.force_zip_64)133        if "r" in mode:134            info = self.info(path)135            out.size = info["size"]136            out.name = info["name"]137        return out138 139    def find(self, path, maxdepth=None, withdirs=False, detail=False, **kwargs):140        if maxdepth is not None and maxdepth < 1:141            raise ValueError("maxdepth must be at least 1")142 143        def to_parts(_path: str):144            return list(filter(None, _path.replace("\\", "/").split("/")))145 146        if not isinstance(path, str):147            path = str(path)148 149        # Remove the leading slash, as the zip file paths are always150        # given without a leading slash151        path = path.lstrip("/")152        path_parts = to_parts(path)153        path_depth = len(path_parts)154 155        self._get_dirs()156 157        result = {}158        # To match posix find, if an exact file name is given, we should159        # return only that file160        if path in self.dir_cache and self.dir_cache[path]["type"] == "file":161            result[path] = self.dir_cache[path]162            return result if detail else [path]163 164        for file_path, file_info in self.dir_cache.items():165            if len(file_parts := to_parts(file_path)) < path_depth or any(166                a != b for a, b in zip(path_parts, file_parts)167            ):168                # skip parent folders and mismatching paths169                continue170 171            if file_info["type"] == "directory":172                if withdirs and file_path not in result:173                    result[file_path.strip("/")] = file_info174                continue175 176            if file_path not in result:177                result[file_path] = file_info if detail else None178 179        if maxdepth:180            result = {181                k: v for k, v in result.items() if k.count("/") < maxdepth + path_depth182            }183        return result if detail else sorted(result)184