|
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": |
|
|
|
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) |
|
|
