Fix execute time not showing up after moving a cell

CHANGELOG.md

@@ -1,3 +1,9 @@
+## 3.1.1 (unreleased)
+
+### Fixed
+
+- Fix execute time not showing up after moving a cell [#109](https://github.com/deshaw/jupyterlab-execute-time/pull/109)
+
 ## [3.1.0](https://github.com/deshaw/jupyterlab-execute-time/compare/v3.0.1...v3.1.0) (2023-11-06)
 
 ### Added

src/ExecuteTimeWidget.ts

@@ -75,21 +75,56 @@ export default class ExecuteTimeWidget extends Widget {
     );
   }
 
+  /**
+   * Handle `CellList.changed` signal.
+   */
   updateConnectedCell(
     sender: CellList,
     changed: IObservableList.IChangedArgs<ICellModel>
   ) {
-    // While we could look at changed.type, it's easier to just remove all
-    // oldValues and add back all new values
+    // When a cell is moved it's model gets re-created so we need to update
+    // the `metadataChanged` listeners.
+
+    // When cells are moved around the `CellList.changed` signal is first
+    // emitted with "add" type and the cell model information and then
+    // with "remove" type but lacking the old model (it is `undefined`).
+    // This causes a problem for the sequence of registering and deregistering
+    // listeners for the `metadataChanged` signal (register can be called when
+    // the cell was no yet removed from `this._cellSlotMap`, and deregister
+    // could be called with `undefined` value hence unable to remove it).
+    // There are two possible solutions:
+    // - (a) go over the list of cells and compare it with `cellSlotMap` (slow)
+    // - (b) deregister the cell model as it gets disposed just before
+    //      `CellList.changed` signals are emitted; we can do this by
+    //       listening to the `ICellModel.sharedModel.disposed` signal.
+    // The (b) solution is implemented in `_registerMetadataChanges` method.
+
+    // Reference:
+    // https://github.com/jupyterlab/jupyterlab/blob/4.0.x/packages/notebook/src/celllist.ts#L131-L159
+
     changed.oldValues.forEach(this._deregisterMetadataChanges.bind(this));
     changed.newValues.forEach(this._registerMetadataChanges.bind(this));
   }
 
   _registerMetadataChanges(cellModel: ICellModel) {
     if (!(cellModel.id in this._cellSlotMap)) {
+      // Register signal handler with `cellModel` stored in closure.
       const fn = () => this._cellMetadataChanged(cellModel);
       this._cellSlotMap[cellModel.id] = fn;
       cellModel.metadataChanged.connect(fn);
+
+      // Copy cell model identifier and store a reference to `metadataChanged`
+      // signal to keep them available even during cell model disposal.
+      const id = cellModel.id;
+      const metadataChanged = cellModel.metadataChanged;
+
+      // Register a model disposal handler on the underlying shared model,
+      // see the explanation in `updateConnectedCell()` method.
+      const deregisterOnDisposal = () => {
+        this._deregisterMetadataChanges({ metadataChanged, id } as ICellModel);
+        cellModel.sharedModel.disposed.disconnect(deregisterOnDisposal);
+      };
+      cellModel.sharedModel.disposed.connect(deregisterOnDisposal);
     }
     // Always re-render cells.
     // In case there was already metadata: do not highlight on first load.

ui-tests/notebooks/Simple_notebook.ipynb

@@ -0,0 +1,55 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e6ba4d4d-2fc5-480d-8f6b-d80e44f51dda",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from time import sleep"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ec671046-8e9f-409e-8b8e-708154c0ac10",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "sleep(0.01)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "000acc20-b975-49b0-905a-b3b11fbfb349",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "sleep(0.01)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.4"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

ui-tests/tests/cell_operations.spec.ts

@@ -0,0 +1,33 @@
+import { expect, galata, test } from '@jupyterlab/galata';
+import { openNotebook, cleanup } from './utils';
+
+const NOTEBOOK_ID = '@jupyterlab/notebook-extension:tracker';
+
+test.describe('Cell operations', () => {
+  test.use({
+    mockSettings: {
+      ...galata.DEFAULT_SETTINGS,
+      [NOTEBOOK_ID]: {
+        ...galata.DEFAULT_SETTINGS[NOTEBOOK_ID],
+        windowingMode: 'defer',
+      },
+    },
+  });
+  test.beforeEach(openNotebook('Simple_notebook.ipynb'));
+  test.afterEach(cleanup);
+
+  test('Add and move a cell', async ({ page }) => {
+    await page.notebook.run();
+    // There are three cells, there should be three widgets
+    expect(await page.locator('.execute-time').count()).toBe(3);
+    // Add a new cell
+    await page.notebook.addCell('code', 'sleep(0.01)');
+    // Move the added cell up (fourth, index three)
+    await page.notebook.selectCells(3);
+    await page.menu.clickMenuItem('Edit>Move Cell Up');
+    // Run the added cell
+    await page.notebook.runCell(2, true);
+    // Four cells should now have the widget
+    expect(await page.locator('.execute-time').count()).toBe(4);
+  });
+});