Implement granular cell updates on Python side

Fix

Author: krassowski

jupyter_ydoc/ynotebook.py

@@ -17,6 +17,8 @@
 # The default minor version of the notebook format.
 NBFORMAT_MINOR_VERSION = 5
 
+_CELL_KEY_TYPE_MAP = {"metadata": Map, "source": Text, "outputs": Array}
+
 
 class YNotebook(YBaseDoc):
     """
@@ -249,7 +251,7 @@ def set(self, value: dict) -> None:
                 "id": str(uuid4()),
             }
         ]
-        old_ycells_by_id = {ycell["id"]: ycell for ycell in self._ycells}
+        old_ycells_by_id: dict[str, Map] = {ycell["id"]: ycell for ycell in self._ycells}
 
         with self._ydoc.transaction():
             new_cell_list: list[dict] = []
@@ -260,7 +262,55 @@ def set(self, value: dict) -> None:
                 cell_id = new_cell.get("id")
                 if cell_id and (old_ycell := old_ycells_by_id.get(cell_id)):
                     old_cell = self._cell_to_py(old_ycell)
-                    if old_cell == new_cell:
+                    updated_granularly = True
+                    if old_cell != new_cell:
+                        # attempt to update cell granularly
+                        old_keys = set(old_cell.keys())
+                        new_keys = set(new_cell.keys())
+
+                        shared_keys = old_keys & new_keys
+                        removed_keys = old_keys - new_keys
+                        added_keys = new_keys - old_keys
+
+                        for key in shared_keys:
+                            if old_cell[key] != new_cell[key]:
+                                if key == "output" and value:
+                                    # outputs require complex handling - some have Text type nested;
+                                    # for now skip creating them; clearing all outputs is fine
+                                    updated_granularly = False
+
+                                if key in _CELL_KEY_TYPE_MAP:
+                                    kind = _CELL_KEY_TYPE_MAP[key]
+                                    value = new_cell[key]
+                                    if kind == Text:
+                                        old: Text = old_ycell[key]
+                                        old.clear()
+                                        old.insert(0, value)
+                                    elif kind == Array:
+                                        old: Array = old_ycell[key]
+                                        old.clear()
+                                        old.extend(value)
+                                    elif kind == Map:
+                                        old: Map = old_ycell[key]
+                                        old.clear()
+                                        for k, v in value.items():
+                                            old[k] = v
+                                else:
+                                    old_ycell[key] = new_cell[key]
+
+                        for key in removed_keys:
+                            del old_ycell[key]
+
+                        for key in added_keys:
+                            if key in _CELL_KEY_TYPE_MAP:
+                                # we hard-reload cells when keys that require nested types get added
+                                # to allow the frontend to connect observers; this could be changed
+                                # in the future, once frontends learn how to observe all changes
+                                updated_granularly = False
+                            else:
+                                old_ycell[key] = new_cell[key]
+
+                    if updated_granularly:
                         new_cell_list.append(old_cell)
                         retained_cells.add(cell_id)
                         continue