Aluode/PerceptionLabPortable

1import errno2import io3import os4import secrets5import shutil6from contextlib import suppress7from functools import cached_property, wraps8from urllib.parse import parse_qs9 10from fsspec.spec import AbstractFileSystem11from fsspec.utils import (12    get_package_version_without_import,13    infer_storage_options,14    mirror_from,15    tokenize,16)17 18 19def wrap_exceptions(func):20    @wraps(func)21    def wrapper(*args, **kwargs):22        try:23            return func(*args, **kwargs)24        except OSError as exception:25            if not exception.args:26                raise27 28            message, *args = exception.args29            if isinstance(message, str) and "does not exist" in message:30                raise FileNotFoundError(errno.ENOENT, message) from exception31            else:32                raise33 34    return wrapper35 36 37PYARROW_VERSION = None38 39 40class ArrowFSWrapper(AbstractFileSystem):41    """FSSpec-compatible wrapper of pyarrow.fs.FileSystem.42 43    Parameters44    ----------45    fs : pyarrow.fs.FileSystem46 47    """48 49    root_marker = "/"50 51    def __init__(self, fs, **kwargs):52        global PYARROW_VERSION53        PYARROW_VERSION = get_package_version_without_import("pyarrow")54        self.fs = fs55        super().__init__(**kwargs)56 57    @property58    def protocol(self):59        return self.fs.type_name60 61    @cached_property62    def fsid(self):63        return "hdfs_" + tokenize(self.fs.host, self.fs.port)64 65    @classmethod66    def _strip_protocol(cls, path):67        ops = infer_storage_options(path)68        path = ops["path"]69        if path.startswith("//"):70            # special case for "hdfs://path" (without the triple slash)71            path = path[1:]72        return path73 74    def ls(self, path, detail=False, **kwargs):75        path = self._strip_protocol(path)76        from pyarrow.fs import FileSelector77 78        try:79            entries = [80                self._make_entry(entry)81                for entry in self.fs.get_file_info(FileSelector(path))82            ]83        except (FileNotFoundError, NotADirectoryError):84            entries = [self.info(path, **kwargs)]85        if detail:86            return entries87        else:88            return [entry["name"] for entry in entries]89 90    def info(self, path, **kwargs):91        path = self._strip_protocol(path)92        [info] = self.fs.get_file_info([path])93        return self._make_entry(info)94 95    def exists(self, path):96        path = self._strip_protocol(path)97        try:98            self.info(path)99        except FileNotFoundError:100            return False101        else:102            return True103 104    def _make_entry(self, info):105        from pyarrow.fs import FileType106 107        if info.type is FileType.Directory:108            kind = "directory"109        elif info.type is FileType.File:110            kind = "file"111        elif info.type is FileType.NotFound:112            raise FileNotFoundError(errno.ENOENT, os.strerror(errno.ENOENT), info.path)113        else:114            kind = "other"115 116        return {117            "name": info.path,118            "size": info.size,119            "type": kind,120            "mtime": info.mtime,121        }122 123    @wrap_exceptions124    def cp_file(self, path1, path2, **kwargs):125        path1 = self._strip_protocol(path1).rstrip("/")126        path2 = self._strip_protocol(path2).rstrip("/")127 128        with self._open(path1, "rb") as lstream:129            tmp_fname = f"{path2}.tmp.{secrets.token_hex(6)}"130            try:131                with self.open(tmp_fname, "wb") as rstream:132                    shutil.copyfileobj(lstream, rstream)133                self.fs.move(tmp_fname, path2)134            except BaseException:135                with suppress(FileNotFoundError):136                    self.fs.delete_file(tmp_fname)137                raise138 139    @wrap_exceptions140    def mv(self, path1, path2, **kwargs):141        path1 = self._strip_protocol(path1).rstrip("/")142        path2 = self._strip_protocol(path2).rstrip("/")143        self.fs.move(path1, path2)144 145    @wrap_exceptions146    def rm_file(self, path):147        path = self._strip_protocol(path)148        self.fs.delete_file(path)149 150    @wrap_exceptions151    def rm(self, path, recursive=False, maxdepth=None):152        path = self._strip_protocol(path).rstrip("/")153        if self.isdir(path):154            if recursive:155                self.fs.delete_dir(path)156            else:157                raise ValueError("Can't delete directories without recursive=False")158        else:159            self.fs.delete_file(path)160 161    @wrap_exceptions162    def _open(self, path, mode="rb", block_size=None, seekable=True, **kwargs):163        if mode == "rb":164            if seekable:165                method = self.fs.open_input_file166            else:167                method = self.fs.open_input_stream168        elif mode == "wb":169            method = self.fs.open_output_stream170        elif mode == "ab":171            method = self.fs.open_append_stream172        else:173            raise ValueError(f"unsupported mode for Arrow filesystem: {mode!r}")174 175        _kwargs = {}176        if mode != "rb" or not seekable:177            if int(PYARROW_VERSION.split(".")[0]) >= 4:178                # disable compression auto-detection179                _kwargs["compression"] = None180        stream = method(path, **_kwargs)181 182        return ArrowFile(self, stream, path, mode, block_size, **kwargs)183 184    @wrap_exceptions185    def mkdir(self, path, create_parents=True, **kwargs):186        path = self._strip_protocol(path)187        if create_parents:188            self.makedirs(path, exist_ok=True)189        else:190            self.fs.create_dir(path, recursive=False)191 192    @wrap_exceptions193    def makedirs(self, path, exist_ok=False):194        path = self._strip_protocol(path)195        self.fs.create_dir(path, recursive=True)196 197    @wrap_exceptions198    def rmdir(self, path):199        path = self._strip_protocol(path)200        self.fs.delete_dir(path)201 202    @wrap_exceptions203    def modified(self, path):204        path = self._strip_protocol(path)205        return self.fs.get_file_info(path).mtime206 207    def cat_file(self, path, start=None, end=None, **kwargs):208        kwargs.setdefault("seekable", start not in [None, 0])209        return super().cat_file(path, start=None, end=None, **kwargs)210 211    def get_file(self, rpath, lpath, **kwargs):212        kwargs.setdefault("seekable", False)213        super().get_file(rpath, lpath, **kwargs)214 215 216@mirror_from(217    "stream",218    [219        "read",220        "seek",221        "tell",222        "write",223        "readable",224        "writable",225        "close",226        "seekable",227    ],228)229class ArrowFile(io.IOBase):230    def __init__(self, fs, stream, path, mode, block_size=None, **kwargs):231        self.path = path232        self.mode = mode233 234        self.fs = fs235        self.stream = stream236 237        self.blocksize = self.block_size = block_size238        self.kwargs = kwargs239 240    def __enter__(self):241        return self242 243    @property244    def size(self):245        return self.stream.size()246 247    def __exit__(self, *args):248        return self.close()249 250 251class HadoopFileSystem(ArrowFSWrapper):252    """A wrapper on top of the pyarrow.fs.HadoopFileSystem253    to connect it's interface with fsspec"""254 255    protocol = "hdfs"256 257    def __init__(258        self,259        host="default",260        port=0,261        user=None,262        kerb_ticket=None,263        replication=3,264        extra_conf=None,265        **kwargs,266    ):267        """268 269        Parameters270        ----------271        host: str272            Hostname, IP or "default" to try to read from Hadoop config273        port: int274            Port to connect on, or default from Hadoop config if 0275        user: str or None276            If given, connect as this username277        kerb_ticket: str or None278            If given, use this ticket for authentication279        replication: int280            set replication factor of file for write operations. default value is 3.281        extra_conf: None or dict282            Passed on to HadoopFileSystem283        """284        from pyarrow.fs import HadoopFileSystem285 286        fs = HadoopFileSystem(287            host=host,288            port=port,289            user=user,290            kerb_ticket=kerb_ticket,291            replication=replication,292            extra_conf=extra_conf,293        )294        super().__init__(fs=fs, **kwargs)295 296    @staticmethod297    def _get_kwargs_from_urls(path):298        ops = infer_storage_options(path)299        out = {}300        if ops.get("host", None):301            out["host"] = ops["host"]302        if ops.get("username", None):303            out["user"] = ops["username"]304        if ops.get("port", None):305            out["port"] = ops["port"]306        if ops.get("url_query", None):307            queries = parse_qs(ops["url_query"])308            if queries.get("replication", None):309                out["replication"] = int(queries["replication"][0])310        return out311