Merge pull request #21 from tomokinakamaru/fix-loader

Fix loader

Author: tomokinakamaru

redspot/datafilters.py

@@ -1,29 +1,21 @@
 from collections import defaultdict
 
-from redspot import database, datafilters
+from redspot import database
 from redspot.notebook import Notebook
 
 
 def load(path):
-    stream = _load(path)
-    stream = datafilters.filter_by_kind(stream)
-    stream = datafilters.filter_by_diff(stream)
-    stream = datafilters.filter_void(stream)
-    yield from stream
-
-
-def _load(path):
     panel_roots = _load_panel_roots(path)
-    for panel, kind, args, notebook in _load_notebooks(path):
+    for time, panel, kind, args, notebook in _load_notebooks(path):
         panel = panel_roots[panel]
-        yield panel, kind, args, notebook
+        yield time, panel, kind, args, notebook
 
 
 def _load_notebooks(path):
     notebooks = defaultdict(Notebook)
-    for _, panel, kind, args in database.get(path):
+    for time, panel, kind, args in database.get(path):
         notebooks[panel].apply(kind, args)
-        yield panel, kind, args, notebooks[panel]
+        yield time, panel, kind, args, notebooks[panel]
 
 
 def _load_panel_roots(path):

redspot/dataloader.py

@@ -10,11 +10,16 @@ def apply(self, kind, args):
         head = kind.split(".", 1)[0]
         self._invoke(head, kind, args)
 
+    def get_cell(self, id):
+        for c in self["cells"]:
+            if c["id"] == id:
+                return c
+
     def _INotebookModel(self, kind, args):
         self._invoke(kind, args)
 
     def _ISharedCell(self, kind, args):
-        cell = self._get_cell(args["cell"])
+        cell = self.get_cell(args["cell"])
         self._invoke(kind, cell, args)
 
     def _INotebookModel_changed_cellsChange(self, args):
@@ -47,11 +52,6 @@ def _invoke(self, key, *args):
         func = getattr(self, f"_{name}", None)
         func and func(*args)
 
-    def _get_cell(self, id):
-        for c in self["cells"]:
-            if c["id"] == id:
-                return c
-
 
 def _apply_delta(obj, args):
     head = 0

redspot/notebook.py

@@ -1,18 +1,94 @@
 from collections import defaultdict
+from copy import deepcopy
 from datetime import datetime
 from itertools import count
 from json import dumps
 
+from nbdime import diff
+
 from redspot import load
 
 
 def main(path, outdir):
     counters = defaultdict(count)
-    for panel, data in load(path):
+    stream = load(path)
+    stream = _filter_by_kind(stream)
+    stream = _filter_by_diff(stream)
+    stream = _filter_void(stream)
+    for panel, data in stream:
         time, uuid = panel.split("-")
         ymdt = datetime.fromtimestamp(int(time) / 1000)
         outd = outdir / datetime.strftime(ymdt, f"%Y-%m%d-%H%M_{uuid}")
         name = f"{next(counters[panel])}.ipynb"
         outd.mkdir(parents=True, exist_ok=True)
         (outd / name).write_text(dumps(data))
     return 0
+
+
+def _filter_void(stream):
+    prev = defaultdict(lambda: None)
+    void = defaultdict(lambda: True)
+    for panel, notebook in stream:
+        if panel in prev:
+            void[panel] = False
+            yield panel, prev[panel]
+        prev[panel] = deepcopy(notebook)
+    for panel, notebook in prev.items():
+        if not void[panel]:
+            yield panel, notebook
+
+
+def _filter_by_diff(stream):
+    notebooks = defaultdict(lambda: {"cells": [], "metadata": {}})
+    for panel, notebook in stream:
+        delta = diff(notebooks[panel], notebook)
+        notebooks[panel] = deepcopy(notebook)
+        if _has_visible_change(delta):
+            yield panel, notebook
+
+
+def _filter_by_kind(stream):
+    notebooks = {}
+    cell_changed = _cell_change_detector()
+    for _, panel, kind, args, notebook in stream:
+        notebooks[panel] = notebook
+        if kind in _yield_immediately:
+            yield panel, notebooks.pop(panel)
+        elif kind in _yield_if_cell_changed:
+            if cell_changed(panel, args):
+                yield panel, notebooks.pop(panel)
+    yield from notebooks.items()
+
+
+def _has_visible_change(delta):
+    for eps in delta:
+        op, key = eps["op"], eps["key"]
+        if op == "patch" and key == "cells":
+            return True
+    return False
+
+
+def _cell_change_detector():
+    def _(panel, args):
+        if panel in prev:
+            new = args.get("cell")
+            old = prev[panel]
+            if old != new:
+                prev[panel] = new
+                return True
+        return False
+
+    prev = {}
+    return _
+
+
+_yield_immediately = (
+    "INotebookModel.changed:cellsChange",
+    "ISharedCell.changed:executionCountChange",
+)
+
+_yield_if_cell_changed = (
+    "ISharedCell.changed:attachmentsChange",
+    "ISharedCell.changed:outputsChange",
+    "ISharedCell.changed:sourceChange",
+)

redspot/replay.py

@@ -13,7 +13,7 @@
 
 def test_load_api():
     g = redspot.load(database)
-    assert len(list(g)) == 3
+    assert len(list(g)) == 16
 
 
 def test_load_cli():