Aluode/PerceptionLabPortable

1import logging2import tarfile3 4import fsspec5from fsspec.archive import AbstractArchiveFileSystem6from fsspec.compression import compr7from fsspec.utils import infer_compression8 9typemap = {b"0": "file", b"5": "directory"}10 11logger = logging.getLogger("tar")12 13 14class TarFileSystem(AbstractArchiveFileSystem):15    """Compressed Tar archives as a file-system (read-only)16 17    Supports the following formats:18    tar.gz, tar.bz2, tar.xz19    """20 21    root_marker = ""22    protocol = "tar"23    cachable = False24 25    def __init__(26        self,27        fo="",28        index_store=None,29        target_options=None,30        target_protocol=None,31        compression=None,32        **kwargs,33    ):34        super().__init__(**kwargs)35        target_options = target_options or {}36 37        if isinstance(fo, str):38            self.of = fsspec.open(fo, protocol=target_protocol, **target_options)39            fo = self.of.open()  # keep the reference40 41        # Try to infer compression.42        if compression is None:43            name = None44 45            # Try different ways to get hold of the filename. `fo` might either46            # be a `fsspec.LocalFileOpener`, an `io.BufferedReader` or an47            # `fsspec.AbstractFileSystem` instance.48            try:49                # Amended io.BufferedReader or similar.50                # This uses a "protocol extension" where original filenames are51                # propagated to archive-like filesystems in order to let them52                # infer the right compression appropriately.53                if hasattr(fo, "original"):54                    name = fo.original55 56                # fsspec.LocalFileOpener57                elif hasattr(fo, "path"):58                    name = fo.path59 60                # io.BufferedReader61                elif hasattr(fo, "name"):62                    name = fo.name63 64                # fsspec.AbstractFileSystem65                elif hasattr(fo, "info"):66                    name = fo.info()["name"]67 68            except Exception as ex:69                logger.warning(70                    f"Unable to determine file name, not inferring compression: {ex}"71                )72 73            if name is not None:74                compression = infer_compression(name)75                logger.info(f"Inferred compression {compression} from file name {name}")76 77        if compression is not None:78            # TODO: tarfile already implements compression with modes like "'r:gz'",79            #  but then would seek to offset in the file work?80            fo = compr[compression](fo)81 82        self._fo_ref = fo83        self.fo = fo  # the whole instance is a context84        self.tar = tarfile.TarFile(fileobj=self.fo)85        self.dir_cache = None86 87        self.index_store = index_store88        self.index = None89        self._index()90 91    def _index(self):92        # TODO: load and set saved index, if exists93        out = {}94        for ti in self.tar:95            info = ti.get_info()96            info["type"] = typemap.get(info["type"], "file")97            name = ti.get_info()["name"].rstrip("/")98            out[name] = (info, ti.offset_data)99 100        self.index = out101        # TODO: save index to self.index_store here, if set102 103    def _get_dirs(self):104        if self.dir_cache is not None:105            return106 107        # This enables ls to get directories as children as well as files108        self.dir_cache = {109            dirname: {"name": dirname, "size": 0, "type": "directory"}110            for dirname in self._all_dirnames(self.tar.getnames())111        }112        for member in self.tar.getmembers():113            info = member.get_info()114            info["name"] = info["name"].rstrip("/")115            info["type"] = typemap.get(info["type"], "file")116            self.dir_cache[info["name"]] = info117 118    def _open(self, path, mode="rb", **kwargs):119        if mode != "rb":120            raise ValueError("Read-only filesystem implementation")121        details, offset = self.index[path]122        if details["type"] != "file":123            raise ValueError("Can only handle regular files")124        return self.tar.extractfile(path)125