Clean up notebook after tidy-imports (#22)

* Clean up notebook after tidy-imports * Update package.json and changelog * Address review comments * Add release date --------- Co-authored-by: Divyansh Choudhary <[email protected]>
deshaw · Feb 20, 2024 · 54333fc · 54333fc
1 parent baeb867
commit 54333fc
Show file tree

Hide file tree

Showing 4 changed files with 1,585 additions and 1,269 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,3 +1,9 @@
+## [5.1.1](https://github.com/deshaw/jupyterlab-pyflyby/compare/v5.1.0...v5.1.1) (2024-02-20)
+
+### Fixed
+
+- Clean up notebook after running tidy-imports
+
 ## [5.1.0](https://github.com/deshaw/jupyterlab-pyflyby/compare/v5.0.0...v5.1.0) (2023-11-21)
 
 ### Added

diff --git a/package.json b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "@deshaw/jupyterlab-pyflyby",
-  "version": "5.1.0",
+  "version": "5.1.1",
   "description": "A labextension to integrate pyflyby with notebooks",
   "keywords": [
     "jupyter",

diff --git a/src/index.tsx b/src/index.tsx
@@ -228,6 +228,7 @@ class PyflyByWidget extends Widget {
     let position = findLinePos(model.cells.get(pyflybyCellIndex));
 
     if (position === -1) {
+      pyflybyCellIndex = 0;
       cell = this._context.model.sharedModel.insertCell(0, {
         source: `${PYFLYBY_START_MSG}\n\n${PYFLYBY_END_MSG}`,
         cell_type: 'code',
@@ -318,7 +319,7 @@ class PyflyByWidget extends Widget {
     for (let i = 0; i < cellArray.length; ++i) {
       const cell = cells.get(i);
       const model = cell.sharedModel;
-      model.setSource(cellArray[i].text);
+      model.setSource(cellArray[i].text.trim());
     }
 
     const joinedImports = imports.trim();
@@ -356,6 +357,39 @@ class PyflyByWidget extends Widget {
           [PYFLYBY_START_MSG, joinedImports, PYFLYBY_END_MSG].join('\n')
         );
     }
+
+    this._cleanupNotebook();
+  }
+
+  _cleanupNotebook() {
+    const cells = this._context.model.cells;
+    const cellsToDelete: number[] = [];
+    for (let i = 0; i < cells.length; ++i) {
+      const cell = cells.get(i);
+      const lines = cell.sharedModel.getSource().split('\n');
+      let shouldDelete = true;
+      for (let j = 0; j < lines.length; ++j) {
+        const line = lines[j].trim();
+        // We are conservative here to only delete the cells which are empty or
+        // just have pyflyby messages
+        if (
+          line !== '' &&
+          line !== PYFLYBY_START_MSG &&
+          line !== PYFLYBY_END_MSG
+        ) {
+          shouldDelete = false;
+        }
+      }
+      if (shouldDelete) {
+        cellsToDelete.push(i);
+      }
+    }
+
+    this._context.model.sharedModel.transact(() => {
+      for (let i = cellsToDelete.length - 1; i >= 0; --i) {
+        this._context.model.sharedModel.deleteCell(cellsToDelete[i]);
+      }
+    });
   }
 
   _fastStringHash(str: string) {