Add ability to trust a notebook

Don Jayamanne (DonJayamanne) · Jul 10, 2020 · Jul 8, 2020 · Jul 8, 2020 · Jul 8, 2020 · Jul 9, 2020
commit 3bbd3fe7c796c2a38349065ee510eabb2ecd7e24
diff --git a/package.json b/package.json
@@ -684,6 +684,25 @@
                 },
                 "enablement": "python.datascience.notebookeditor.canrestartNotebookkernel"
             },
+            {
+                "command": "python.datascience.notebookeditor.trust",
+                "title": "%DataScience.trustNotebookCommandTitle%",
+                "category": "Python",
+                "icon": {
+                    "light": "resources/light/un-trusted.svg",
+                    "dark": "resources/dark/un-trusted.svg"
+                },
+                "enablement": "notebookEditorFocused && !python.datascience.isnotebooktrusted && python.datascience.trustfeatureenabled"
+            },
+            {
+                "command": "python.datascience.notebookeditor.trusted",
+                "title": "%DataScience.notebookIsTrusted%",
+                "category": "Python",
+                "icon": {
+                    "light": "resources/light/trusted.svg",
+                    "dark": "resources/dark/trusted.svg"
+                }
+            },
             {
                 "command": "python.datascience.notebookeditor.runallcells",
                 "title": "%python.command.python.datascience.notebookeditor.runallcells.title%",
@@ -849,6 +868,18 @@
                     "title": "%python.command.python.datascience.restartkernel.title%",
                     "group": "navigation",
                     "when": "notebookEditorFocused"
+                },
+                {
+                    "command": "python.datascience.notebookeditor.trust",
+                    "title": "%DataScience.trustNotebookCommandTitle%",
+                    "group": "navigation@1",
+                    "when": "notebookEditorFocused && !python.datascience.isnotebooktrusted && python.datascience.trustfeatureenabled"
+                },
+                {
+                    "command": "python.datascience.notebookeditor.trusted",
+                    "title": "%DataScience.notebookIsTrusted%",
+                    "group": "navigation@1",
+                    "when": "notebookEditorFocused && python.datascience.isnotebooktrusted && python.datascience.trustfeatureenabled"
                 }
             ],
             "explorer/context": [
@@ -1122,6 +1153,18 @@
                     "category": "Python",
                     "when": "python.datascience.isnativeactive && python.datascience.featureenabled && python.datascience.isnotebooktrusted"
                 },
+                {
+                    "command": "python.datascience.notebookeditor.trust",
+                    "title": "%DataScience.trustNotebookCommandTitle%",
+                    "category": "Python",
+                    "when": "python.datascience.featureenabled && notebookEditorFocused && !python.datascience.isnotebooktrusted && python.datascience.trustfeatureenabled"
+                },
+                {
+                    "command": "python.datascience.notebookeditor.trusted",
+                    "title": "%DataScience.notebookIsTrusted%",
+                    "category": "Python",
+                    "when": "config.noExists"
+                },
                 {
                     "command": "python.datascience.notebookeditor.runallcells",
                     "title": "%python.command.python.datascience.notebookeditor.runallcells.title%",

diff --git a/package.nls.json b/package.nls.json
@@ -370,6 +370,7 @@
     "DataScience.fetchingDataViewer": "Fetching data ...",
     "DataScience.noRowsInDataViewer": "No rows match current filter",
     "DataScience.jupyterServer": "Jupyter Server",
+    "DataScience.trustNotebookCommandTitle": "Trust notebook",
     "DataScience.notebookIsTrusted": "Trusted",
     "DataScience.notebookIsNotTrusted": "Not Trusted",
     "DataScience.noKernel": "No Kernel",

diff --git a/resources/dark/trusted.svg b/resources/dark/trusted.svg
diff --git a/resources/dark/un-trusted.svg b/resources/dark/un-trusted.svg
diff --git a/resources/light/run-file.svg b/resources/light/run-file.svg
diff --git a/resources/light/trusted.svg b/resources/light/trusted.svg
diff --git a/resources/light/un-trusted.svg b/resources/light/un-trusted.svg
diff --git a/src/client/common/application/commands.ts b/src/client/common/application/commands.ts
@@ -183,4 +183,6 @@ export interface ICommandNameArgumentTypeMapping extends ICommandNameWithoutArgu
     [DSCommands.OpenNotebookNonCustomEditor]: [Uri];
     [DSCommands.GatherQuality]: [string];
     [DSCommands.EnableLoadingWidgetsFrom3rdPartySource]: [undefined | never];
+    [DSCommands.TrustNotebook]: [undefined | never | Uri];
+    [DSCommands.TrustedNotebook]: [undefined | never];
 }
diff --git a/src/client/datascience/constants.ts b/src/client/datascience/constants.ts
@@ -90,6 +90,8 @@ export namespace Commands {
     export const SaveAsNotebookNonCustomEditor = 'python.datascience.notebookeditor.saveAs';
     export const OpenNotebookNonCustomEditor = 'python.datascience.notebookeditor.open';
     export const GatherQuality = 'python.datascience.gatherquality';
+    export const TrustNotebook = 'python.datascience.notebookeditor.trust';
+    export const TrustedNotebook = 'python.datascience.notebookeditor.trusted';
     export const EnableLoadingWidgetsFrom3rdPartySource =
         'python.datascience.enableLoadingWidgetScriptsFromThirdPartySource';
 }

diff --git a/src/client/datascience/context/activeEditorContext.ts b/src/client/datascience/context/activeEditorContext.ts
@@ -122,7 +122,7 @@ export class ActiveEditorContextService implements IExtensionSingleActivationSer
                 .getOrCreateNotebook({ identity: activeEditor.file, getOnly: true })
                 .then((nb) => {
                     if (activeEditor === this.notebookEditorProvider.activeEditor) {
-                        const canStart = nb && nb.status !== ServerStatus.NotStarted;
+                        const canStart = nb && nb.status !== ServerStatus.NotStarted && activeEditor.model?.isTrusted;
                         this.canRestartNotebookKernelContext.set(!!canStart).ignoreErrors();
                     }
                 })

diff --git a/src/client/datascience/interactive-ipynb/nativeEditor.ts b/src/client/datascience/interactive-ipynb/nativeEditor.ts
@@ -6,7 +6,6 @@ import * as path from 'path';
 import {
     CancellationToken,
     CancellationTokenSource,
-    commands,
     Event,
     EventEmitter,
     Memento,
@@ -98,7 +97,6 @@ import { KernelSwitcher } from '../jupyter/kernels/kernelSwitcher';
 const nativeEditorDir = path.join(EXTENSION_ROOT_DIR, 'out', 'datascience-ui', 'notebook');
 @injectable()
 export class NativeEditor extends InteractiveBase implements INotebookEditor {
-    public readonly type: 'old' | 'custom' = 'custom';
     public get onDidChangeViewState(): Event<void> {
         return this._onDidChangeViewState.event;
     }
@@ -140,6 +138,7 @@ export class NativeEditor extends InteractiveBase implements INotebookEditor {
     public get isDirty(): boolean {
         return this.model ? this.model.isDirty : false;
     }
+    public readonly type: 'old' | 'custom' = 'custom';
     public model: Readonly<INotebookModel> | undefined;
     protected savedEvent: EventEmitter<INotebookEditor> = new EventEmitter<INotebookEditor>();
     protected closedEvent: EventEmitter<INotebookEditor> = new EventEmitter<INotebookEditor>();
@@ -152,6 +151,7 @@ export class NativeEditor extends InteractiveBase implements INotebookEditor {
     private startupTimer: StopWatch = new StopWatch();
     private loadedAllCells: boolean = false;
     private executeCancelTokens = new Set<CancellationTokenSource>();
+    private previouslyNotTrusted:boolean = false;
 
     constructor(
         @multiInject(IInteractiveWindowListener) listeners: IInteractiveWindowListener[],
@@ -227,6 +227,7 @@ export class NativeEditor extends InteractiveBase implements INotebookEditor {
         );
         asyncRegistry.push(this);
 
+        asyncRegistry.push(this.trustService.onDidSetNotebookTrust(this.monitorChangesToTrust, this));
         this.synchronizer.subscribeToUserActions(this, this.postMessage.bind(this));
     }
 
@@ -251,6 +252,7 @@ export class NativeEditor extends InteractiveBase implements INotebookEditor {
 
         // Sign up for dirty events
         model.changed(this.modelChanged.bind(this));
+        this.previouslyNotTrusted = model.isTrusted;
     }
 
     // tslint:disable-next-line: no-any
@@ -596,7 +598,13 @@ export class NativeEditor extends InteractiveBase implements INotebookEditor {
             }
         }
     }
-
+    private async monitorChangesToTrust() {
+        if (this.previouslyNotTrusted && this.model?.isTrusted) {
+            this.previouslyNotTrusted = false;
+            // Tell UI to update main state
+            this.postMessage(InteractiveWindowMessages.TrustNotebookComplete).ignoreErrors();
+        }
+    }
     private renameVariableExplorerHeights(name: string, updatedName: string) {
         // Updates the workspace storage to reflect the updated name of the notebook
         // should be called if the name of the notebook changes
@@ -612,38 +620,8 @@ export class NativeEditor extends InteractiveBase implements INotebookEditor {
     }
 
     private async launchNotebookTrustPrompt() {
-        const prompts = [
-            localize.DataScience.trustNotebook(),
-            localize.DataScience.doNotTrustNotebook(),
-            localize.DataScience.trustAllNotebooks()
-        ];
-        const selection = await this.applicationShell.showErrorMessage(
-            localize.DataScience.launchNotebookTrustPrompt(),
-            ...prompts
-        );
-        if (!selection) {
-            return;
-        }
-        if (this.model && selection === localize.DataScience.trustNotebook() && !this.model.isTrusted) {
-            try {
-                const contents = this.model.getContent();
-                await this.trustService.trustNotebook(this.model.file, contents);
-                // Update model trust
-                this.model.update({
-                    source: 'user',
-                    kind: 'updateTrust',
-                    oldDirty: this.model.isDirty,
-                    newDirty: this.model.isDirty,
-                    isNotebookTrusted: true
-                });
-                // Tell UI to update main state
-                await this.postMessage(InteractiveWindowMessages.TrustNotebookComplete);
-            } catch (err) {
-                traceError(err);
-            }
-        } else if (selection === localize.DataScience.trustAllNotebooks()) {
-            // Take the user to the settings UI where they can manually turn on the alwaysTrustNotebooks setting
-            commands.executeCommand('workbench.action.openSettings', 'python.dataScience.alwaysTrustNotebooks');
+        if (this.model && !this.model.isTrusted) {
+            await this.commandManager.executeCommand(Commands.TrustNotebook, this.model.file);
         }
     }
 
@@ -831,4 +809,4 @@ export class NativeEditor extends InteractiveBase implements INotebookEditor {
             }
         }
     }
-}
+
diff --git a/src/client/datascience/notebook/helpers/cellUpdateHelpers.ts b/src/client/datascience/notebook/helpers/cellUpdateHelpers.ts
@@ -20,8 +20,13 @@ import { traceError } from '../../../common/logger';
 import { sendTelemetryEvent } from '../../../telemetry';
 import { VSCodeNativeTelemetry } from '../../constants';
 import { VSCodeNotebookModel } from '../../notebookStorage/vscNotebookModel';
+import { INotebookModel } from '../../types';
 import { findMappedNotebookCellModel } from './cellMappers';
-import { createCellFromVSCNotebookCell, updateVSCNotebookCellMetadata } from './helpers';
+import {
+    createCellFromVSCNotebookCell,
+    createVSCCellOutputsFromOutputs,
+    updateVSCNotebookCellMetadata
+} from './helpers';
 // tslint:disable-next-line: no-var-requires no-require-imports
 const vscodeNotebookEnums = require('vscode') as typeof import('vscode-proposed');
 
@@ -53,9 +58,8 @@ export function updateCellModelWithChangesToVSCCell(
 
 /**
  * Updates a notebook document as a result of trusting it.
- * @returns `true` if document has been updated, else `false`.
  */
-export function updateVSCNotebookAfterTrustingNotebook(document: NotebookDocument) {
+export function updateVSCNotebookAfterTrustingNotebook(document: NotebookDocument, model: INotebookModel) {
     const areAllCellsEditableAndRunnable = document.cells.every((cell) => {
         if (cell.cellKind === vscodeNotebookEnums.CellKind.Markdown) {
             return cell.metadata.editable;
@@ -69,8 +73,9 @@ export function updateVSCNotebookAfterTrustingNotebook(document: NotebookDocumen
         document.metadata.editable &&
         document.metadata.runnable;
 
+    // If already trusted, then nothing to do.
     if (isDocumentEditableAndRunnable && areAllCellsEditableAndRunnable) {
-        return false;
+        return;
     }
 
     document.metadata.cellEditable = true;
@@ -83,8 +88,14 @@ export function updateVSCNotebookAfterTrustingNotebook(document: NotebookDocumen
         if (cell.cellKind !== vscodeNotebookEnums.CellKind.Markdown) {
             cell.metadata.runnable = true;
         }
+
+        // Restore the output once we trust the notebook.
+        const cellModel = findMappedNotebookCellModel(cell, model.cells);
+        if (cellModel) {
+            // tslint:disable-next-line: no-any
+            cell.outputs = createVSCCellOutputsFromOutputs(cellModel.data.outputs as any);
+        }
     });
-    return true;
 }
 
 export function clearCellForExecution(vscCell: NotebookCell, model: VSCodeNotebookModel) {