Introduced mimeTypeHandler and executionSummary and refactored code

subhash-arabhi · Achal1607 · commit d1db013b77b4 · 2025-09-12T18:25:35.000+05:30
diff --git a/vscode/src/notebooks/codeCellExecution.ts b/vscode/src/notebooks/codeCellExecution.ts
@@ -1,9 +1,10 @@
 import { commands, NotebookCell, NotebookCellExecution, NotebookCellOutput, NotebookController } from "vscode";
 import { LOGGER } from "../logger";
 import { NotebookCellExecutionResult } from "../lsp/protocol";
-import { createErrorOutputItem, createOutputItem } from "./utils";
+import { createErrorOutputItem } from "./utils";
 import { nbCommands } from "../commands/commands";
 import { mimeTypes } from "./constants";
+import { MimeTypeHandler } from "./mimeTypeHandler";
 
 export class CodeCellExecution {
     private controller?: NotebookController;
@@ -74,7 +75,7 @@ export class CodeCellExecution {
             await this.execution!.replaceOutputItems(createErrorOutputItem(updatedData), this.output);
         } else {
             await this.execution!.replaceOutputItems(
-                createOutputItem(updatedData, mimeType),
+                new MimeTypeHandler(mimeType).makeOutputItem(updatedData),
                 this.output
             );
         }
diff --git a/vscode/src/notebooks/executionSummary.ts b/vscode/src/notebooks/executionSummary.ts
@@ -0,0 +1,51 @@
+/*
+ * Copyright (c) 2024, Oracle and/or its affiliates.
+ *
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+export interface ExecutionSummaryData {
+  executionOrder?: number | null;
+  success?: boolean;
+}
+
+export class ExecutionSummary {
+  constructor(
+    public executionOrder: number | null = null,
+    public success: boolean = false
+  ) {}
+
+  static fromMetadata(
+    meta?: ExecutionSummaryData,
+    fallbackExecCount?: number | null
+  ): ExecutionSummary {
+    const order =
+      meta?.executionOrder != null
+        ? meta.executionOrder
+        : fallbackExecCount ?? null;
+    const success = meta?.success ?? false;
+    return new ExecutionSummary(order, success);
+  }
+
+  toJSON(): ExecutionSummaryData {
+    return {
+      executionOrder: this.executionOrder,
+      success: this.success,
+    };
+  }
+}
diff --git a/vscode/src/notebooks/kernel.ts b/vscode/src/notebooks/kernel.ts
@@ -22,15 +22,16 @@
 import { globalState } from '../globalState';
 import { isNbCommandRegistered } from '../commands/utils';
 import { nbCommands } from '../commands/commands';
-import { createErrorOutput, createOutputItem } from './utils';
 import { NotebookCellExecutionResult } from '../lsp/protocol';
 import { NotebookCell, NotebookController, NotebookDocument, Disposable, notebooks, commands, NotebookCellOutput, workspace } from 'vscode';
 import { LanguageClient } from 'vscode-languageclient/node';
 import { l10n } from '../localiser';
 import { ijnbConstants, ipynbConstants, supportLanguages } from './constants';
 import { LOGGER } from '../logger';
 import { CodeCellExecution } from './codeCellExecution';
-import { isError } from '../utils';
+import { isError, isString } from '../utils';
+import { MimeTypeHandler } from './mimeTypeHandler';
+import { createErrorOutput } from './utils';
 
 export class IJNBKernel implements Disposable {
   private readonly controllers: NotebookController[] = [];
@@ -165,7 +166,7 @@ export class IJNBKernel implements Disposable {
     try {
       exec.start(Date.now());
       await exec.replaceOutput([
-        new NotebookCellOutput([createOutputItem(cell.document.getText(), mimeType)]),
+        new NotebookCellOutput([new MimeTypeHandler(mimeType).makeOutputItem(cell.document.getText())]),
       ]);
       exec.end(true, Date.now());
     } catch (error) {
@@ -175,10 +176,10 @@ export class IJNBKernel implements Disposable {
   }
 
   cleanUpKernel = workspace.onDidCloseNotebookDocument(doc => {
-  if (doc.notebookType === ijnbConstants.NOTEBOOK_TYPE) {
-    IJNBKernel.executionCounter.delete(doc.uri.toString());
-  }
-});
+    if (doc.notebookType === ijnbConstants.NOTEBOOK_TYPE) {
+      IJNBKernel.executionCounter.delete(doc.uri.toString());
+    }
+  });
 }
 
 export const notebookKernel = new IJNBKernel();
diff --git a/vscode/src/notebooks/mimeTypeHandler.ts b/vscode/src/notebooks/mimeTypeHandler.ts
@@ -0,0 +1,71 @@
+/*
+ * Copyright (c) 2024, Oracle and/or its affiliates.
+ *
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+import { Buffer } from 'buffer';
+import * as vscode from 'vscode';
+import { IMimeBundle } from './types';
+import { mimeTypes } from './constants';
+
+export type DataOrBytes = string | Uint8Array;
+
+export class MimeTypeHandler {
+  constructor(public readonly value: string) {}
+  get isText(): boolean {
+    return this.value === mimeTypes.TEXT;
+  }
+  get isImage(): boolean {
+    return this.value.startsWith('image/');
+  }
+
+  static toBytes(data: DataOrBytes): Uint8Array {
+    if (typeof data === 'string') {
+      return Buffer.from(data, 'base64');
+    }
+    return data;
+  }
+
+  static toString(data: DataOrBytes): string {
+    if (typeof data === 'string') {
+      return data;
+    }
+    return new TextDecoder().decode(data);
+  }
+
+    makeOutputItem(data: DataOrBytes): vscode.NotebookCellOutputItem {
+    if (this.isImage) {
+      const bytes = MimeTypeHandler.toBytes(data);
+      return new vscode.NotebookCellOutputItem(bytes, this.value);
+    }
+    const text = MimeTypeHandler.toString(data);
+    return vscode.NotebookCellOutputItem.text(text, this.value);
+  }
+
+    static itemsFromBundle(bundle: IMimeBundle): vscode.NotebookCellOutputItem[] {
+    return Object.entries(bundle).flatMap(([mime, data]) => {
+      const mt = new MimeTypeHandler(mime);
+      if (mt.isText || mt.isImage) {
+        const payload = Array.isArray(data) ? data.join('') : data;
+        return [mt.makeOutputItem(payload)];
+      }
+      return [];
+    });
+  }
+}
diff --git a/vscode/src/notebooks/utils.ts b/vscode/src/notebooks/utils.ts
@@ -33,6 +33,9 @@ import {
 import { randomUUID } from 'crypto';
 import { isString } from '../utils';
 import { mimeTypes } from './constants';
+import { MimeTypeHandler } from './mimeTypeHandler';
+import { ExecutionSummary } from './executionSummary';
+
 
 export function base64ToUint8Array(base64: string): Uint8Array {
   if (typeof Buffer !== 'undefined' && typeof Buffer.from === 'function') {
@@ -56,15 +59,6 @@ export function uint8ArrayToBase64(data: Uint8Array): string {
   return btoa(binary);
 }
 
-export function createOutputItem(data: string | Uint8Array, mimeType: string): vscode.NotebookCellOutputItem {
-  if (mimeType.startsWith('image/')) {
-    const bytes = typeof data === 'string' ? base64ToUint8Array(data) : data;
-    return new vscode.NotebookCellOutputItem(bytes, mimeType);
-  }
-  const text = typeof data === 'string' ? data : new TextDecoder().decode(data);
-  return vscode.NotebookCellOutputItem.text(text, mimeType);
-}
-
 export const createErrorOutput = (err: string | Error) => {
   return new vscode.NotebookCellOutput([createErrorOutputItem(err)]);
 }
@@ -86,14 +80,13 @@ export function parseCell(cell: ICell): vscode.NotebookCellData {
   const cellData = new vscode.NotebookCellData(kind, value, language);
   cellData.metadata = { id: cell.id, ...cell.metadata };
   if (cell.cell_type === 'code') {
-
-    const metaExec = (cell.metadata as IMetadata).executionSummary;
-    const executionOrder = metaExec?.executionOrder ?? cell.execution_count ?? undefined;
-    const success = metaExec?.success ?? undefined;
-
+    const execSummary = ExecutionSummary.fromMetadata(
+      (cell.metadata as IMetadata).executionSummary,
+      cell.execution_count ?? null
+    );
     cellData.executionSummary = {
-      executionOrder,
-      success,
+      executionOrder: execSummary.executionOrder ?? undefined,
+      success: execSummary.success,
     };
 
     if (Array.isArray(cell.outputs)) {
@@ -119,7 +112,7 @@ export function parseOutput(raw: IOutput): vscode.NotebookCellOutput[] {
     case 'stream':
       outputs.push(
         new vscode.NotebookCellOutput([
-          vscode.NotebookCellOutputItem.text(Array.isArray(raw.text) ? raw.text.join('') : raw.text),
+          new MimeTypeHandler(mimeTypes.TEXT).makeOutputItem(Array.isArray(raw.text) ? raw.text.join('') : raw.text),
         ])
       );
       break;
@@ -138,19 +131,8 @@ export function parseOutput(raw: IOutput): vscode.NotebookCellOutput[] {
 
     case 'display_data':
     case 'execute_result':
-      const items: vscode.NotebookCellOutputItem[] = [];
-      const bundle = raw.data || {};
-      for (const mime in bundle) {
-        const data = bundle[mime];
-        if (mime === mimeTypes.TEXT) {
-          const text = Array.isArray(data) ? data.join('') : String(data);
-          items.push(vscode.NotebookCellOutputItem.text(text));
-        } else if ((mime as string).startsWith('image/')) {
-          const b64 = Array.isArray(data) ? data.join('') : String(data);
-          const bytes = base64ToUint8Array(b64);
-          items.push(new vscode.NotebookCellOutputItem(bytes, mime));
-        }
-      }
+      const bundle = raw.data ?? {};
+      const items = MimeTypeHandler.itemsFromBundle(bundle);
       if (items.length) {
         outputs.push(new vscode.NotebookCellOutput(items, raw.metadata));
       }
@@ -168,7 +150,8 @@ export function serializeCell(cell: vscode.NotebookCellData): ICell {
     const executionCount = exec.executionOrder ?? null;
     const success = exec.success ?? false;
 
-    const metadata = { ...baseMeta, executionSummary: { executionOrder: executionCount, success } };
+    const execSummary = new ExecutionSummary(executionCount, success);
+    const metadata = { ...baseMeta, executionSummary: execSummary.toJSON()};
 
     const outputs: IOutput[] = (cell.outputs || []).map((output): IOutput => {
       const data: IMimeBundle = {};
