Aluode/PerceptionLabPortable

1import base642import io3import re4 5import requests6 7import fsspec8 9 10class JupyterFileSystem(fsspec.AbstractFileSystem):11    """View of the files as seen by a Jupyter server (notebook or lab)"""12 13    protocol = ("jupyter", "jlab")14 15    def __init__(self, url, tok=None, **kwargs):16        """17 18        Parameters19        ----------20        url : str21            Base URL of the server, like "http://127.0.0.1:8888". May include22            token in the string, which is given by the process when starting up23        tok : str24            If the token is obtained separately, can be given here25        kwargs26        """27        if "?" in url:28            if tok is None:29                try:30                    tok = re.findall("token=([a-z0-9]+)", url)[0]31                except IndexError as e:32                    raise ValueError("Could not determine token") from e33            url = url.split("?", 1)[0]34        self.url = url.rstrip("/") + "/api/contents"35        self.session = requests.Session()36        if tok:37            self.session.headers["Authorization"] = f"token {tok}"38 39        super().__init__(**kwargs)40 41    def ls(self, path, detail=True, **kwargs):42        path = self._strip_protocol(path)43        r = self.session.get(f"{self.url}/{path}")44        if r.status_code == 404:45            raise FileNotFoundError(path)46        r.raise_for_status()47        out = r.json()48 49        if out["type"] == "directory":50            out = out["content"]51        else:52            out = [out]53        for o in out:54            o["name"] = o.pop("path")55            o.pop("content")56            if o["type"] == "notebook":57                o["type"] = "file"58        if detail:59            return out60        return [o["name"] for o in out]61 62    def cat_file(self, path, start=None, end=None, **kwargs):63        path = self._strip_protocol(path)64        r = self.session.get(f"{self.url}/{path}")65        if r.status_code == 404:66            raise FileNotFoundError(path)67        r.raise_for_status()68        out = r.json()69        if out["format"] == "text":70            # data should be binary71            b = out["content"].encode()72        else:73            b = base64.b64decode(out["content"])74        return b[start:end]75 76    def pipe_file(self, path, value, **_):77        path = self._strip_protocol(path)78        json = {79            "name": path.rsplit("/", 1)[-1],80            "path": path,81            "size": len(value),82            "content": base64.b64encode(value).decode(),83            "format": "base64",84            "type": "file",85        }86        self.session.put(f"{self.url}/{path}", json=json)87 88    def mkdir(self, path, create_parents=True, **kwargs):89        path = self._strip_protocol(path)90        if create_parents and "/" in path:91            self.mkdir(path.rsplit("/", 1)[0], True)92        json = {93            "name": path.rsplit("/", 1)[-1],94            "path": path,95            "size": None,96            "content": None,97            "type": "directory",98        }99        self.session.put(f"{self.url}/{path}", json=json)100 101    def mv(self, path1, path2, recursive=False, maxdepth=None, **kwargs):102        if path1 == path2:103            return104        self.session.patch(f"{self.url}/{path1}", json={"path": path2})105 106    def _rm(self, path):107        path = self._strip_protocol(path)108        self.session.delete(f"{self.url}/{path}")109 110    def _open(self, path, mode="rb", **kwargs):111        path = self._strip_protocol(path)112        if mode == "rb":113            data = self.cat_file(path)114            return io.BytesIO(data)115        else:116            return SimpleFileWriter(self, path, mode="wb")117 118 119class SimpleFileWriter(fsspec.spec.AbstractBufferedFile):120    def _upload_chunk(self, final=False):121        """Never uploads a chunk until file is done122 123        Not suitable for large files124        """125        if final is False:126            return False127        self.buffer.seek(0)128        data = self.buffer.read()129        self.fs.pipe_file(self.path, data)130