Adding the outputs manager.

Author: ellisonbg

jupyter_rtc_core/outputs/manager.py

@@ -0,0 +1,100 @@
+import json
+import os
+from pathlib import Path, PurePath
+import shutil
+
+
+from traitlets.config.configurable import LoggingConfigurable
+from traitlets import (
+    Any,
+    Bool,
+    Dict,
+    Instance,
+    List,
+    TraitError,
+    Type,
+    Unicode,
+    default,
+    validate,
+)
+
+from jupyter_core.paths import jupyter_runtime_dir
+
+class OutputsManager(LoggingConfigurable):
+
+    outputs_path = Instance(PurePath, help="The local runtime dir")
+    _last_output_index = Dict(default_value={})
+
+    @default("outputs_path")
+    def _default_outputs_path(self):
+        return Path(jupyter_runtime_dir()) / "outputs"
+    
+    def _ensure_path(self, file_id, cell_id):
+        nested_dir = self.outputs_path / file_id / cell_id
+        self.log.info(f"Creating directory: {nested_dir}")
+        nested_dir.mkdir(parents=True, exist_ok=True)
+
+    def _build_path(self, file_id, cell_id=None, output_index=None):
+        path = self.outputs_path / file_id
+        if cell_id is not None:
+            path = path / cell_id
+        if output_index is not None:
+            path = path / f"{output_index}.output"
+        return path
+    
+    def get(self, file_id, cell_id, output_index):
+        """Get an outputs by file_id, cell_id, and output_index."""
+        self.log.info(f"OutputsManager.get: {file_id} {cell_id} {output_index}")
+        path = self._build_path(file_id, cell_id, output_index)
+        if not os.path.isfile(path):
+            raise FileNotFoundError(f"The output file doesn't exist: {path}")
+        with open(path, "r", encoding="utf-8") as f:
+            output = json.loads(f.read())
+        return output
+
+    def get_stream(self, file_id, cell_id):
+        "Get the stream output for a cell by file_id and cell_id."
+        path = self._build_path(file_id, cell_id) / "stream"
+        if not os.path.isfile(path):
+            raise FileNotFoundError(f"The output file doesn't exist: {path}")
+        with open(path, "r", encoding="utf-8") as f:
+            output = f.read()
+        return output
+
+    def write(self, file_id, cell_id, output):
+        """Write a new output for file_id and cell_id."""
+        self.log.info(f"OutputsManager.write: {file_id} {cell_id} {output}")
+        if output["output_type"] == "stream":
+            url = self.write_stream(file_id, cell_id, output)
+        else:
+            url = self.write_output(file_id, cell_id, output)
+        return url
+
+    def write_output(self, file_id, cell_id, output):
+        self._ensure_path(file_id, cell_id)
+        last_index = self._last_output_index.get(cell_id, -1)
+        index = last_index + 1
+        self._last_output_index[cell_id] = index
+        path = self._build_path(file_id, cell_id, index)
+        data = json.dumps(output, ensure_ascii=False)
+        with open(path, "w", encoding="utf-8") as f:
+            f.write(data)
+        url = f"/api/outputs/{file_id}/{cell_id}/{index}.output"
+        return url
+
+    def write_stream(self, file_id, cell_id, output):
+        self._ensure_path(file_id, cell_id)
+        path = self._build_path(file_id, cell_id) / "stream"
+        text = output["text"]
+        mode = 'a' if os.path.isfile(path) else 'w'
+        with open(path, "a", encoding="utf-8") as f:
+            f.write(text)
+        url = f"/api/outputs/{file_id}/{cell_id}/stream"
+        return url
+    
+    def clear(self, file_id, cell_id=None):
+        path = self._build_path(file_id, cell_id)
+        shutil.rmtree(path)
+
+
+