Fix bugs due to pyflyby cell (#20)

* Fix bugs due to pyflyby cell

* Add changelog and update version

* Address review comments

* Fix typo

* Fix function documentation

---------

Co-authored-by: Divyansh Choudhary <choudhdi@deshaw.com>

CHANGELOG.md

@@ -10,6 +10,15 @@ Support to run `tidy-imports` on a notebook.
 
 - **Breaking**: Ported to JupyterLab 4.x
 
+## [4.3.2](https://github.com/deshaw/jupyterlab-pyflyby/compare/v4.3.1...v4.3.2) (UNRELEASED)
+
+### Fixed
+
+- Fixed a bug where multiple clicks of the TidyImportsButton caused the notebook to be modified in unexpected ways. Some of these were -
+  - Outputs of cells being mismatched when reconstructing the notebook.
+  - Addition of random cells above and below the pyflyby tagged cell.
+  - Mutliple clicks of the TidyImportsButton would keep adding extra lines in the pyflyby cell.
+
 ## [4.3.1](https://github.com/deshaw/jupyterlab-pyflyby/compare/v4.3.0...v4.3.1) (2023-08-10)
 
 ### Fixed

package.json

@@ -1,6 +1,6 @@
 {
   "name": "@deshaw/jupyterlab-pyflyby",
-  "version": "4.3.1",
+  "version": "4.3.2",
   "description": "A labextension to integrate pyflyby with notebooks",
   "keywords": [
     "jupyter",

src/cellUtils.ts

@@ -1,7 +1,7 @@
 import { toArray } from '@lumino/algorithm';
 import { MultilineString } from '@jupyterlab/nbformat';
 import { ICellModel } from '@jupyterlab/cells';
-import { PYFLYBY_END_MSG } from './constants';
+import { PYFLYBY_END_MSG, PYFLYBY_START_MSG } from './constants';
 
 // FIXME: There's got to be a better Typescript solution
 // for distinguishing between members of a union type at runtime.
@@ -119,3 +119,32 @@ export const findLinePos = (cell: ICellModel): number => {
   // These imports will be moved to next cell
   return -1;
 };
+
+/**
+ * This code extracts non-imports lines from the pyflyby cell.
+ *
+ * @param cell - a cell model
+ */
+export const extractCodeFromPyflybyCell = (cell: ICellModel): string => {
+  const lines: string[] = normalizeMultilineString(cell.toJSON().source);
+
+  let stIdx = -1,
+    enIdx = -1;
+  for (let i = 0; i < lines.length; ++i) {
+    if (lines[i] === PYFLYBY_START_MSG.trim()) {
+      stIdx = i;
+    }
+    if (lines[i] === PYFLYBY_END_MSG.trim()) {
+      enIdx = i;
+    }
+  }
+
+  // we splice it twice to remove the pyflyby messages
+  const imports: string = lines
+    .splice(stIdx, enIdx - stIdx + 1)
+    .splice(stIdx + 1, enIdx - stIdx - 1)
+    .join();
+  const remainingCode: string = lines.join('');
+  console.log(imports, remainingCode);
+  return remainingCode;
+};

src/index.tsx

@@ -19,7 +19,11 @@
 
 // Lumino imports
 import { toArray, ArrayExt } from '@lumino/algorithm';
-import { JSONValue, JSONObject } from '@lumino/coreutils';
+import {
+  JSONValue,
+  JSONObject,
+  ReadonlyPartialJSONValue
+} from '@lumino/coreutils';
 import { Widget, Panel } from '@lumino/widgets';
 
 // Jupyterlab imports
@@ -45,7 +49,7 @@ import { debug } from 'debug';
 import React from 'react';
 
 // relative imports
-import { findCell, findLinePos } from './cellUtils';
+import { extractCodeFromPyflybyCell, findCell, findLinePos } from './cellUtils';
 import {
   PYFLYBY_CELL_TAG,
   PYFLYBY_START_MSG,
@@ -307,27 +311,61 @@ class PyflyByWidget extends Widget {
   }
 
   restoreNotebookAfterTidyImports(cellArray: any, imports: any): void {
-    const { cellIndex } = this._findAndSetImportCoordinates();
     const cells = this._context.model.cells;
     for (let i = 0; i < cellArray.length; ++i) {
       const cell = cells.get(i);
       cell.value.remove(0, cell.value.text.length);
       cell.value.insert(0, cellArray[i].text.trim());
     }
-    const joined_imports = imports.join('\n').trim();
-    if (cells.get(0).value.text.length === 0) {
-      cells.get(0).value.insert(0, joined_imports);
-    } else {
+    const joinedImports = imports.trim();
+    const { cellIndex } = this._findAndSetImportCoordinates();
+
+    const remainingCodeInPyflybyCell = extractCodeFromPyflybyCell(
+      cells.get(cellIndex)
+    );
+
+    if (remainingCodeInPyflybyCell !== '') {
+      // Remove import statements from the current pyflyby cell
+      cells
+        .get(cellIndex)
+        .value.remove(0, cells.get(cellIndex).value.text.length);
+
+      cells.get(cellIndex).value.insert(0, remainingCodeInPyflybyCell);
+
+      // remove the pyflyby cell tag from the current cell, we'll create another pyflyby cell
+      const tags: string[] = cells
+        .get(cellIndex)
+        .metadata.get('tags') as string[];
+      tags.splice(tags.indexOf(PYFLYBY_CELL_TAG));
+      cells
+        .get(cellIndex)
+        .metadata.set('tags', tags as ReadonlyPartialJSONValue);
+
+      // Create a new pyflyby cell and insert it at the top
       const cell = this._context.model.contentFactory.createCodeCell({
         cell: {
-          source: joined_imports,
+          source: [PYFLYBY_START_MSG, joinedImports, PYFLYBY_END_MSG].join(
+            '\n'
+          ),
           cell_type: 'code',
           metadata: {
-            trusted: true
+            trusted: true,
+            tags: [PYFLYBY_CELL_TAG]
           }
         }
       });
-      cells.insert(cellIndex, cell);
+
+      cells.insert(0, cell);
+    } else {
+      cells
+        .get(cellIndex)
+        .value.remove(0, cells.get(cellIndex).value.text.length);
+      cells
+        .get(cellIndex)
+        .value.insert(
+          0,
+          [PYFLYBY_START_MSG, joinedImports, PYFLYBY_END_MSG].join('\n')
+        );
     }
   }