airflow_scheduler/.venv/Lib/site-packages/fsspec/implementations/jupyter.py

import base64
import io
import re

import requests

import fsspec


class JupyterFileSystem(fsspec.AbstractFileSystem):
    """View of the files as seen by a Jupyter server (notebook or lab)"""

    protocol = ("jupyter", "jlab")

    def __init__(self, url, tok=None, **kwargs):
        """

        Parameters
        ----------
        url : str
            Base URL of the server, like "http://127.0.0.1:8888". May include
            token in the string, which is given by the process when starting up
        tok : str
            If the token is obtained separately, can be given here
        kwargs
        """
        if "?" in url:
            if tok is None:
                try:
                    tok = re.findall("token=([a-z0-9]+)", url)[0]
                except IndexError as e:
                    raise ValueError("Could not determine token") from e
            url = url.split("?", 1)[0]
        self.url = url.rstrip("/") + "/api/contents"
        self.session = requests.Session()
        if tok:
            self.session.headers["Authorization"] = f"token {tok}"

        super().__init__(**kwargs)

    def ls(self, path, detail=True, **kwargs):
        path = self._strip_protocol(path)
        r = self.session.get(f"{self.url}/{path}")
        if r.status_code == 404:
            return FileNotFoundError(path)
        r.raise_for_status()
        out = r.json()

        if out["type"] == "directory":
            out = out["content"]
        else:
            out = [out]
        for o in out:
            o["name"] = o.pop("path")
            o.pop("content")
            if o["type"] == "notebook":
                o["type"] = "file"
        if detail:
            return out
        return [o["name"] for o in out]

    def cat_file(self, path, start=None, end=None, **kwargs):
        path = self._strip_protocol(path)
        r = self.session.get(f"{self.url}/{path}")
        if r.status_code == 404:
            return FileNotFoundError(path)
        r.raise_for_status()
        out = r.json()
        if out["format"] == "text":
            # data should be binary
            b = out["content"].encode()
        else:
            b = base64.b64decode(out["content"])
        return b[start:end]

    def pipe_file(self, path, value, **_):
        path = self._strip_protocol(path)
        json = {
            "name": path.rsplit("/", 1)[-1],
            "path": path,
            "size": len(value),
            "content": base64.b64encode(value).decode(),
            "format": "base64",
            "type": "file",
        }
        self.session.put(f"{self.url}/{path}", json=json)

    def mkdir(self, path, create_parents=True, **kwargs):
        path = self._strip_protocol(path)
        if create_parents and "/" in path:
            self.mkdir(path.rsplit("/", 1)[0], True)
        json = {
            "name": path.rsplit("/", 1)[-1],
            "path": path,
            "size": None,
            "content": None,
            "type": "directory",
        }
        self.session.put(f"{self.url}/{path}", json=json)

    def _rm(self, path):
        path = self._strip_protocol(path)
        self.session.delete(f"{self.url}/{path}")

    def _open(self, path, mode="rb", **kwargs):
        path = self._strip_protocol(path)
        if mode == "rb":
            data = self.cat_file(path)
            return io.BytesIO(data)
        else:
            return SimpleFileWriter(self, path, mode="wb")


class SimpleFileWriter(fsspec.spec.AbstractBufferedFile):
    def _upload_chunk(self, final=False):
        """Never uploads a chunk until file is done

        Not suitable for large files
        """
        if final is False:
            return False
        self.buffer.seek(0)
        data = self.buffer.read()
        self.fs.pipe_file(self.path, data)