Centralize accessing nb metadata (#15389)

* Centralize accessing nb metadata * oops
microsoft · Mar 18, 2024 · ffeaa79 · ffeaa79
1 parent 076cb1c
commit ffeaa79
Show file tree

Hide file tree

Showing 7 changed files with 45 additions and 47 deletions.
diff --git a/src/interactive-window/interactiveWindow.ts b/src/interactive-window/interactiveWindow.ts
@@ -1,7 +1,6 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
-import type * as nbformat from '@jupyterlab/nbformat';
 import * as path from '../platform/vscode-path/path';
 import {
     Event,
@@ -54,7 +53,6 @@ import {
 } from './editor-integration/types';
 import { IDataScienceErrorHandler } from '../kernels/errors/types';
 import { CellExecutionCreator } from '../kernels/execution/cellExecutionCreator';
-import { updateNotebookMetadata } from '../kernels/execution/helpers';
 import { chainWithPendingUpdates } from '../kernels/execution/notebookUpdater';
 import { generateMarkdownFromCodeLines, parseForComments } from '../platform/common/utils';
 import { KernelController } from '../kernels/kernelController';
@@ -606,12 +604,6 @@ export class InteractiveWindow implements IInteractiveWindow {
         }
         const kernel = this.controller.kernel?.value;
 
-        // Pull out the metadata from our active notebook
-        const metadata: nbformat.INotebookMetadata = {};
-        if (kernel) {
-            await updateNotebookMetadata(metadata, kernel.kernelConnectionMetadata);
-        }
-
         let defaultFileName;
         if (this.submitters && this.submitters.length) {
             const lastSubmitter = this.submitters[this.submitters.length - 1];

diff --git a/src/kernels/execution/cellExecution.ts b/src/kernels/execution/cellExecution.ts
@@ -38,6 +38,7 @@ import { isKernelSessionDead } from '../kernel';
 import { ICellExecution } from './types';
 import { KernelError } from '../errors/kernelError';
 import { getCachedSysPrefix } from '../../platform/interpreter/helpers';
+import { getCellMetadata } from '../../platform/common/utils/jupyter';
 
 /**
  * Factory for CellExecution objects.
@@ -416,7 +417,7 @@ export class CellExecution implements ICellExecution, IDisposable {
         // eslint-disable-next-line @typescript-eslint/no-explicit-any
         const metadata: any = {
             ...{ cellId: this.cell.document.uri.toString() },
-            ...(this.cell.metadata?.custom?.metadata || {}) // Send the Cell Metadata
+            ...getCellMetadata(this.cell).metadata // Send the Cell Metadata
         };
 
         const kernelConnection = session.kernel;

diff --git a/src/kernels/execution/helpers.ts b/src/kernels/execution/helpers.ts
@@ -655,7 +655,7 @@ export function hasErrorOutput(outputs: readonly NotebookCellOutput[]) {
 }
 
 // eslint-disable-next-line complexity
-export async function updateNotebookMetadata(
+export async function updateNotebookMetadataWithSelectedKernel(
     metadata?: nbformat.INotebookMetadata,
     kernelConnection?: KernelConnectionMetadata,
     kernelInfo?: Partial<KernelMessage.IInfoReplyMsg['content']>

diff --git a/src/kernels/execution/helpers.unit.test.ts b/src/kernels/execution/helpers.unit.test.ts
@@ -5,7 +5,11 @@ import * as sinon from 'sinon';
 import type * as nbformat from '@jupyterlab/nbformat';
 import { assert } from 'chai';
 import { Uri } from 'vscode';
-import { cellOutputToVSCCellOutput, getNotebookCellOutputMetadata, updateNotebookMetadata } from './helpers';
+import {
+    cellOutputToVSCCellOutput,
+    getNotebookCellOutputMetadata,
+    updateNotebookMetadataWithSelectedKernel
+} from './helpers';
 import { IJupyterKernelSpec, PythonKernelConnectionMetadata } from '../types';
 import { PythonEnvironment } from '../../platform/pythonEnvironments/info';
 import { PythonExtension } from '@vscode/python-extension';
@@ -56,12 +60,12 @@ suite(`UpdateNotebookMetadata`, () => {
         disposables = dispose(disposables);
     });
     test('Empty call does not change anything', async () => {
-        const value = await updateNotebookMetadata();
+        const value = await updateNotebookMetadataWithSelectedKernel();
         assert.strictEqual(value.changed, false);
     });
     test('Ensure Language', async () => {
         const notebookMetadata = { orig_nbformat: 4 };
-        const value = await updateNotebookMetadata(notebookMetadata);
+        const value = await updateNotebookMetadataWithSelectedKernel(notebookMetadata);
 
         // Verify lang info added
         verifyMetadata(notebookMetadata, { orig_nbformat: 4, language_info: { name: '' } });
@@ -74,7 +78,7 @@ suite(`UpdateNotebookMetadata`, () => {
             interpreter: python36Global,
             kernelSpec: pythonDefaultKernelSpec
         });
-        const value = await updateNotebookMetadata(notebookMetadata, kernelConnection);
+        const value = await updateNotebookMetadataWithSelectedKernel(notebookMetadata, kernelConnection);
 
         // Verify lang info added
         verifyMetadata(notebookMetadata, {
@@ -92,7 +96,7 @@ suite(`UpdateNotebookMetadata`, () => {
             interpreter: python37Global,
             kernelSpec: pythonDefaultKernelSpec
         });
-        const value = await updateNotebookMetadata(notebookMetadata, kernelConnection);
+        const value = await updateNotebookMetadataWithSelectedKernel(notebookMetadata, kernelConnection);
 
         // Verify version updated 3.6 => 3.7
         verifyMetadata(notebookMetadata, {
@@ -114,7 +118,7 @@ suite(`UpdateNotebookMetadata`, () => {
             interpreter: python36Global,
             kernelSpec: pythonDefaultKernelSpec
         });
-        const value = await updateNotebookMetadata(notebookMetadata, kernelConnection);
+        const value = await updateNotebookMetadataWithSelectedKernel(notebookMetadata, kernelConnection);
 
         // Verify kernel_spec name updated JUNK => python3
         verifyMetadata(notebookMetadata, {
@@ -146,7 +150,7 @@ suite(`UpdateNotebookMetadata`, () => {
             interpreter: python36Global,
             kernelSpec: pythonDefaultKernelSpec
         });
-        const value = await updateNotebookMetadata(notebookMetadata, kernelConnection);
+        const value = await updateNotebookMetadataWithSelectedKernel(notebookMetadata, kernelConnection);
 
         // Verify display_name updated due to interpreter hash change
         verifyMetadata(notebookMetadata, {
@@ -194,7 +198,7 @@ suite(`UpdateNotebookMetadata`, () => {
             interpreter: python36Global,
             kernelSpec: pythonDefaultKernelSpec
         });
-        const value = await updateNotebookMetadata(notebookMetadata, kernelConnection);
+        const value = await updateNotebookMetadataWithSelectedKernel(notebookMetadata, kernelConnection);
 
         // Verify display_name updated due to interpreter hash change
         verifyMetadata(newNotebookMetadata, {

diff --git a/src/notebooks/controllers/vscodeNotebookController.ts b/src/notebooks/controllers/vscodeNotebookController.ts
@@ -1,7 +1,6 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
-import type * as nbformat from '@jupyterlab/nbformat';
 import {
     CancellationError,
     CancellationTokenSource,
@@ -15,14 +14,12 @@ import {
     NotebookCellKind,
     NotebookController,
     NotebookDocument,
-    NotebookEdit,
     NotebookEditor,
     NotebookRendererScript,
     notebooks,
     Uri,
     window,
-    workspace,
-    WorkspaceEdit
+    workspace
 } from 'vscode';
 import { IPythonExtensionChecker } from '../../platform/api/types';
 import { Exiting, InteractiveWindowView, JupyterNotebookView, PYTHON_LANGUAGE } from '../../platform/common/constants';
@@ -59,7 +56,7 @@ import {
 } from '../../kernels/types';
 import { KernelDeadError } from '../../kernels/errors/kernelDeadError';
 import { DisplayOptions } from '../../kernels/displayOptions';
-import { getNotebookMetadata, isJupyterNotebook } from '../../platform/common/utils';
+import { getNotebookMetadata, isJupyterNotebook, updateNotebookMetadata } from '../../platform/common/utils';
 import { ConsoleForegroundColors } from '../../platform/logging/types';
 import { KernelConnector } from './kernelConnector';
 import { IConnectionDisplayData, IConnectionDisplayDataProvider, IVSCodeNotebookController } from './types';
@@ -68,7 +65,7 @@ import { CellExecutionCreator } from '../../kernels/execution/cellExecutionCreat
 import {
     traceCellMessage,
     endCellAndDisplayErrorsInCell,
-    updateNotebookMetadata
+    updateNotebookMetadataWithSelectedKernel
 } from '../../kernels/execution/helpers';
 import type { KernelMessage } from '@jupyterlab/services';
 import { initializeInteractiveOrNotebookTelemetryBasedOnUserAction } from '../../kernels/telemetry/helper';
@@ -722,28 +719,10 @@ async function updateNotebookDocumentMetadata(
     kernelConnection?: KernelConnectionMetadata,
     kernelInfo?: Partial<KernelMessage.IInfoReplyMsg['content']>
 ) {
-    let metadata = getNotebookMetadata(document) || {};
-    const { changed } = await updateNotebookMetadata(metadata, kernelConnection, kernelInfo);
+    const metadata = getNotebookMetadata(document) || {};
+    const { changed } = await updateNotebookMetadataWithSelectedKernel(metadata, kernelConnection, kernelInfo);
     if (changed) {
-        const edit = new WorkspaceEdit();
-        // Create a clone.
-        const docMetadata = JSON.parse(
-            JSON.stringify(
-                (document.metadata as {
-                    custom?: Exclude<Partial<nbformat.INotebookContent>, 'cells'>;
-                }) || { custom: {} }
-            )
-        );
-
-        docMetadata.custom = docMetadata.custom || {};
-        docMetadata.custom.metadata = metadata;
-        edit.set(document.uri, [
-            NotebookEdit.updateNotebookMetadata({
-                ...(document.metadata || {}),
-                custom: docMetadata.custom
-            })
-        ]);
-        await workspace.applyEdit(edit);
+        await updateNotebookMetadata(document, metadata);
     }
 }
 

diff --git a/src/notebooks/export/exportUtil.ts b/src/notebooks/export/exportUtil.ts
@@ -32,7 +32,7 @@ export abstract class ExportUtilBase implements IExportUtil {
             return data;
         });
         const notebookData = new NotebookData(cellData);
-        notebookData.metadata = document.metadata;
+        notebookData.metadata = JSON.parse(JSON.stringify(document.metadata));
         return serializerApi.exports.exportNotebook(notebookData);
     }
 

diff --git a/src/platform/common/utils.ts b/src/platform/common/utils.ts
@@ -4,7 +4,7 @@
 import { SemVer, parse } from 'semver';
 import type * as nbformat from '@jupyterlab/nbformat';
 import * as uriPath from '../../platform/vscode-path/resources';
-import { NotebookData, NotebookDocument, TextDocument, Uri, workspace } from 'vscode';
+import { NotebookData, NotebookDocument, NotebookEdit, TextDocument, Uri, WorkspaceEdit, workspace } from 'vscode';
 import {
     InteractiveWindowView,
     jupyterLanguageToMonacoLanguageMapping,
@@ -177,6 +177,28 @@ export function getNotebookFormat(document: NotebookDocument): {
     };
 }
 
+export async function updateNotebookMetadata(document: NotebookDocument, metadata: NotebookMetadata) {
+    const edit = new WorkspaceEdit();
+    // Create a clone.
+    const docMetadata = JSON.parse(
+        JSON.stringify(
+            (document.metadata as {
+                custom?: Exclude<Partial<nbformat.INotebookContent>, 'cells'>;
+            }) || { custom: {} }
+        )
+    );
+
+    docMetadata.custom = docMetadata.custom || {};
+    docMetadata.custom.metadata = metadata;
+    edit.set(document.uri, [
+        NotebookEdit.updateNotebookMetadata({
+            ...(document.metadata || {}),
+            custom: docMetadata.custom
+        })
+    ]);
+    await workspace.applyEdit(edit);
+}
+
 export function getAssociatedJupyterNotebook(document: TextDocument): NotebookDocument | undefined {
     return workspace.notebookDocuments.find(
         (notebook) => isJupyterNotebook(notebook) && notebook.getCells().some((cell) => cell.document === document)