Store isInteractiveWindowMessageCell at root level

microsoft · Mar 18, 2024 · b97a4fe · b97a4fe
1 parent 52429e6
commit b97a4fe
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 10 deletions.
diff --git a/src/interactive-window/systemInfoCell.ts b/src/interactive-window/systemInfoCell.ts
@@ -16,6 +16,7 @@ import { getDisplayNameOrNameOfKernelConnection } from '../kernels/helpers';
 import { SysInfoReason } from '../messageTypes';
 import { DataScience } from '../platform/common/utils/localize';
 import { KernelConnectionMetadata } from '../kernels/types';
+import { getCellMetadata } from '../platform/common/utils/jupyter';
 
 export function getStartConnectMessage(kernelMetadata: KernelConnectionMetadata, reason: SysInfoReason) {
     const displayName = getDisplayNameOrNameOfKernelConnection(kernelMetadata);
@@ -38,11 +39,7 @@ export function getFinishConnectMessage(kernelMetadata: KernelConnectionMetadata
 }
 
 export function isSysInfoCell(cell: NotebookCell) {
-    return (
-        cell.kind === NotebookCellKind.Markup &&
-        cell.metadata.custom?.metadata &&
-        cell.metadata.custom.metadata['isInteractiveWindowMessageCell']
-    );
+    return cell.kind === NotebookCellKind.Markup && getCellMetadata(cell).isInteractiveWindowMessageCell === true;
 }
 
 export class SystemInfoCell {
@@ -70,7 +67,7 @@ export class SystemInfoCell {
         let addedCellIndex: number | undefined;
         await chainWithPendingUpdates(this.notebookDocument, (edit) => {
             const markdownCell = new NotebookCellData(NotebookCellKind.Markup, message, MARKDOWN_LANGUAGE);
-            markdownCell.metadata = { custom: { metadata: { isInteractiveWindowMessageCell: true } } };
+            markdownCell.metadata = { isInteractiveWindowMessageCell: true };
             addedCellIndex = this.notebookDocument.cellCount;
             const nbEdit = NotebookEdit.insertCells(addedCellIndex, [markdownCell]);
             edit.set(this.notebookDocument.uri, [nbEdit]);
@@ -87,9 +84,7 @@ export class SystemInfoCell {
                     edit.replace(cell.document.uri, new Range(0, 0, cell.document.lineCount, 0), newMessage);
 
                     edit.set(this.notebookDocument!.uri, [
-                        NotebookEdit.updateCellMetadata(cell.index, {
-                            custom: { metadata: { isInteractiveWindowMessageCell: true } }
-                        })
+                        NotebookEdit.updateCellMetadata(cell.index, { isInteractiveWindowMessageCell: true })
                     ]);
                     return;
                 }

diff --git a/src/notebooks/export/exportToPythonPlain.ts b/src/notebooks/export/exportToPythonPlain.ts
@@ -8,6 +8,7 @@ import { IConfigurationService } from '../../platform/common/types';
 import { appendLineFeed } from '../../platform/common/utils';
 import { IExport } from './types';
 import { ServiceContainer } from '../../platform/ioc/container';
+import { getCellMetadata } from '../../platform/common/utils/jupyter';
 
 // Handles exporting a NotebookDocument to python without using nbconvert
 export class ExportToPythonPlain implements IExport {
@@ -42,7 +43,7 @@ export class ExportToPythonPlain implements IExport {
     private exportDocument(document: NotebookDocument): string {
         return document
             .getCells()
-            .filter((cell) => !cell.metadata.custom?.metadata?.isInteractiveWindowMessageCell) // We don't want interactive window sys info cells
+            .filter((cell) => !getCellMetadata(cell).isInteractiveWindowMessageCell) // We don't want interactive window sys info cells
             .reduce((previousValue, currentValue) => previousValue + this.exportCell(currentValue), '');
     }
 

diff --git a/src/platform/common/utils/jupyter.ts b/src/platform/common/utils/jupyter.ts
@@ -0,0 +1,22 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+import { NotebookCellKind, type NotebookCell, type NotebookCellData } from 'vscode';
+import type * as nbformat from '@jupyterlab/nbformat';
+
+type JupyterCellMetadata = Pick<nbformat.IRawCell, 'id' | 'metadata' | 'attachments'> &
+    Pick<nbformat.IMarkdownCell, 'id' | 'attachments'> &
+    Pick<nbformat.ICodeCell, 'id' | 'metadata' | 'attachments'> &
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    Record<string, any>;
+
+export function getCellMetadata(cell: NotebookCell | NotebookCellData): JupyterCellMetadata {
+    const metadata: JupyterCellMetadata = cell.metadata?.custom || {};
+    if (cell.kind !== NotebookCellKind.Markup) {
+        const cellMetadata = metadata as nbformat.IRawCell;
+        // metadata property is never optional.
+        cellMetadata.metadata = cellMetadata.metadata || {};
+    }
+
+    return metadata;
+}