diff --git a/build/ci/vscode-python-pr-validation.yaml b/build/ci/vscode-python-pr-validation.yaml index 1d60dd550c82..f46d58cbef18 100644 --- a/build/ci/vscode-python-pr-validation.yaml +++ b/build/ci/vscode-python-pr-validation.yaml @@ -56,6 +56,7 @@ stages: 'DataScience': TestsToRun: 'testDataScience' NeedsPythonTestReqs: true + NeedsPythonFunctionalReqs: true 'Smoke': TestsToRun: 'testSmoke' NeedsPythonTestReqs: true diff --git a/package.json b/package.json index 09d8c0370ddd..6504ceee0878 100644 --- a/package.json +++ b/package.json @@ -694,6 +694,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%", @@ -860,6 +879,18 @@ "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" + }, { "command": "python.datascience.export", "title": "%DataScience.notebookExportAs%", @@ -1138,6 +1169,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 index 244a4ffefc9b..b8941cffecd4 100644 --- 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 new file mode 100644 index 000000000000..a841903129b0 --- /dev/null +++ b/resources/dark/trusted.svg @@ -0,0 +1,4 @@ + + + + diff --git a/resources/dark/un-trusted.svg b/resources/dark/un-trusted.svg new file mode 100644 index 000000000000..f97f5f561d9e --- /dev/null +++ b/resources/dark/un-trusted.svg @@ -0,0 +1,4 @@ + + + diff --git a/resources/light/run-file.svg b/resources/light/run-file.svg index f41a5c8fa2d8..0adc9e411b03 100644 --- a/resources/light/run-file.svg +++ b/resources/light/run-file.svg @@ -1,3 +1,4 @@ - - + + diff --git a/resources/light/trusted.svg b/resources/light/trusted.svg new file mode 100644 index 000000000000..82e41b9ff47b --- /dev/null +++ b/resources/light/trusted.svg @@ -0,0 +1,4 @@ + + + diff --git a/resources/light/un-trusted.svg b/resources/light/un-trusted.svg new file mode 100644 index 000000000000..c9be7cd88dda --- /dev/null +++ b/resources/light/un-trusted.svg @@ -0,0 +1,4 @@ + + + diff --git a/src/client/common/application/commands.ts b/src/client/common/application/commands.ts index 17c73a6e3a43..dc310105fdb1 100644 --- 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 index 7967ba05ce8b..7d7b732a1e4c 100644 --- 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 index 39d7ad76cdcf..103021f29481 100644 --- 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/export/exportUtil.ts b/src/client/datascience/export/exportUtil.ts index d7fe774d9923..516e42b52c75 100644 --- a/src/client/datascience/export/exportUtil.ts +++ b/src/client/datascience/export/exportUtil.ts @@ -63,7 +63,7 @@ export class ExportUtil { await this.jupyterExporter.exportToFile(cells, tempFile.filePath, false); const newPath = path.join(tempDir.path, '.ipynb'); await this.fileSystem.copyFile(tempFile.filePath, newPath); - model = await this.notebookStorage.load(Uri.file(newPath)); + model = await this.notebookStorage.get(Uri.file(newPath)); } finally { tempFile.dispose(); tempDir.dispose(); @@ -73,7 +73,7 @@ export class ExportUtil { } public async removeSvgs(source: Uri) { - const model = await this.notebookStorage.load(source); + const model = await this.notebookStorage.get(source); const newCells: ICell[] = []; for (const cell of model.cells) { diff --git a/src/client/datascience/interactive-ipynb/nativeEditor.ts b/src/client/datascience/interactive-ipynb/nativeEditor.ts index c40a543ad097..0ae442b19afb 100644 --- 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 { 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 | undefined; protected savedEvent: EventEmitter = new EventEmitter(); protected closedEvent: EventEmitter = new EventEmitter(); @@ -152,6 +151,7 @@ export class NativeEditor extends InteractiveBase implements INotebookEditor { private startupTimer: StopWatch = new StopWatch(); private loadedAllCells: boolean = false; private executeCancelTokens = new Set(); + 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); } } diff --git a/src/client/datascience/interactive-ipynb/nativeEditorProvider.ts b/src/client/datascience/interactive-ipynb/nativeEditorProvider.ts index 53e1a39506b7..2b0445be4cf5 100644 --- a/src/client/datascience/interactive-ipynb/nativeEditorProvider.ts +++ b/src/client/datascience/interactive-ipynb/nativeEditorProvider.ts @@ -184,7 +184,7 @@ export class NativeEditorProvider implements INotebookEditorProvider, CustomEdit this.untitledCounter = getNextUntitledCounter(file, this.untitledCounter); // Load our model from our storage object. - const model = await this.storage.load(file, contents, options); + const model = await this.storage.get(file, contents, options); // Make sure to listen to events on the model this.trackModel(model); diff --git a/src/client/datascience/interactive-ipynb/nativeEditorStorage.ts b/src/client/datascience/interactive-ipynb/nativeEditorStorage.ts index 0586697ef17b..10f2e29eb8ed 100644 --- a/src/client/datascience/interactive-ipynb/nativeEditorStorage.ts +++ b/src/client/datascience/interactive-ipynb/nativeEditorStorage.ts @@ -74,20 +74,20 @@ export class NativeEditorStorage implements INotebookStorage { return `${path.basename(model.file.fsPath)}-${uuid()}`; } - public load( + public get( file: Uri, possibleContents?: string, backupId?: string, forVSCodeNotebook?: boolean ): Promise; - public load( + public get( file: Uri, possibleContents?: string, // tslint:disable-next-line: unified-signatures skipDirtyContents?: boolean, forVSCodeNotebook?: boolean ): Promise; - public load( + public get( file: Uri, possibleContents?: string, // tslint:disable-next-line: no-any diff --git a/src/client/datascience/interactive-ipynb/notebookStorageProvider.ts b/src/client/datascience/interactive-ipynb/notebookStorageProvider.ts index 0461c7e5f6ea..5ffb2814f4f4 100644 --- a/src/client/datascience/interactive-ipynb/notebookStorageProvider.ts +++ b/src/client/datascience/interactive-ipynb/notebookStorageProvider.ts @@ -56,8 +56,8 @@ export class NotebookStorageProvider implements INotebookStorageProvider { public deleteBackup(model: INotebookModel, backupId?: string) { return this.storage.deleteBackup(model, backupId); } - public load(file: Uri, contents?: string, backupId?: string, forVSCodeNotebook?: boolean): Promise; - public load( + public get(file: Uri, contents?: string, backupId?: string, forVSCodeNotebook?: boolean): Promise; + public get( file: Uri, contents?: string, // tslint:disable-next-line: unified-signatures @@ -66,7 +66,7 @@ export class NotebookStorageProvider implements INotebookStorageProvider { ): Promise; // tslint:disable-next-line: no-any - public load(file: Uri, contents?: string, options?: any, forVSCodeNotebook?: boolean): Promise { + public get(file: Uri, contents?: string, options?: any, forVSCodeNotebook?: boolean): Promise { const key = file.toString(); if (!this.storageAndModels.has(key)) { // Every time we load a new untitled file, up the counter past the max value for this counter @@ -74,7 +74,7 @@ export class NotebookStorageProvider implements INotebookStorageProvider { file, NotebookStorageProvider.untitledCounter ); - const promise = this.storage.load(file, contents, options, forVSCodeNotebook); + const promise = this.storage.get(file, contents, options, forVSCodeNotebook); this.storageAndModels.set(key, promise.then(this.trackModel.bind(this))); } return this.storageAndModels.get(key)!; @@ -90,7 +90,7 @@ export class NotebookStorageProvider implements INotebookStorageProvider { const uri = this.getNextNewNotebookUri(forVSCodeNotebooks); // Always skip loading from the hot exit file. When creating a new file we want a new file. - return this.load(uri, contents, true); + return this.get(uri, contents, true); } private getNextNewNotebookUri(forVSCodeNotebooks?: boolean): Uri { diff --git a/src/client/datascience/interactive-ipynb/trustCommandHandler.ts b/src/client/datascience/interactive-ipynb/trustCommandHandler.ts new file mode 100644 index 000000000000..09ee31f49e62 --- /dev/null +++ b/src/client/datascience/interactive-ipynb/trustCommandHandler.ts @@ -0,0 +1,76 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +'use strict'; + +import { inject, injectable } from 'inversify'; +import { Uri } from 'vscode'; +import { IExtensionSingleActivationService } from '../../activation/types'; +import { IApplicationShell, ICommandManager } from '../../common/application/types'; +import { ContextKey } from '../../common/contextKey'; +import { EnableTrustedNotebooks } from '../../common/experiments/groups'; +import '../../common/extensions'; +import { IDisposableRegistry, IExperimentService } from '../../common/types'; +import { swallowExceptions } from '../../common/utils/decorators'; +import { DataScience } from '../../common/utils/localize'; +import { noop } from '../../common/utils/misc'; +import { Commands } from '../constants'; +import { INotebookStorageProvider } from '../interactive-ipynb/notebookStorageProvider'; +import { INotebookEditorProvider, ITrustService } from '../types'; + +@injectable() +export class TrustCommandHandler implements IExtensionSingleActivationService { + constructor( + @inject(ITrustService) private readonly trustService: ITrustService, + @inject(INotebookEditorProvider) private readonly editorProvider: INotebookEditorProvider, + @inject(INotebookStorageProvider) private readonly storageProvider: INotebookStorageProvider, + @inject(ICommandManager) private readonly commandManager: ICommandManager, + @inject(IApplicationShell) private readonly applicationShell: IApplicationShell, + @inject(IDisposableRegistry) private readonly disposables: IDisposableRegistry, + @inject(IExperimentService) private readonly experiments: IExperimentService + ) {} + public async activate(): Promise { + this.activateInBackground().ignoreErrors(); + } + public async activateInBackground(): Promise { + if (!(await this.experiments.inExperiment(EnableTrustedNotebooks.experiment))) { + return; + } + const context = new ContextKey('python.datascience.trustfeatureenabled', this.commandManager); + context.set(true).ignoreErrors(); + this.disposables.push(this.commandManager.registerCommand(Commands.TrustNotebook, this.onTrustNotebook, this)); + this.disposables.push(this.commandManager.registerCommand(Commands.TrustedNotebook, noop)); + } + @swallowExceptions('Trusting notebook') + private async onTrustNotebook(uri?: Uri) { + uri = uri ?? this.editorProvider.activeEditor?.file; + if (!uri) { + return; + } + + const model = await this.storageProvider.get(uri); + if (model.isTrusted) { + return; + } + + const selection = await this.applicationShell.showErrorMessage( + DataScience.launchNotebookTrustPrompt(), + DataScience.trustNotebook(), + DataScience.doNotTrustNotebook(), + DataScience.trustAllNotebooks() + ); + if (selection !== DataScience.trustNotebook() || model.isTrusted) { + return; + } + // Update model trust + model.update({ + source: 'user', + kind: 'updateTrust', + oldDirty: model.isDirty, + newDirty: model.isDirty, + isNotebookTrusted: true + }); + const contents = model.getContent(); + await this.trustService.trustNotebook(model.file, contents); + } +} diff --git a/src/client/datascience/notebook/contentProvider.ts b/src/client/datascience/notebook/contentProvider.ts index 3e48f4c88e84..622298bf7cd1 100644 --- a/src/client/datascience/notebook/contentProvider.ts +++ b/src/client/datascience/notebook/contentProvider.ts @@ -72,8 +72,8 @@ export class NotebookContentProvider implements INotebookContentProvider { } // If there's no backup id, then skip loading dirty contents. const model = await (openContext.backupId - ? this.notebookStorage.load(uri, undefined, openContext.backupId, true) - : this.notebookStorage.load(uri, undefined, true, true)); + ? this.notebookStorage.get(uri, undefined, openContext.backupId, true) + : this.notebookStorage.get(uri, undefined, true, true)); setSharedProperty('ds_notebookeditor', 'native'); sendTelemetryEvent(Telemetry.CellCount, undefined, { count: model.cells.length }); @@ -81,7 +81,7 @@ export class NotebookContentProvider implements INotebookContentProvider { } @captureTelemetry(Telemetry.Save, undefined, true) public async saveNotebook(document: NotebookDocument, cancellation: CancellationToken) { - const model = await this.notebookStorage.load(document.uri, undefined, undefined, true); + const model = await this.notebookStorage.get(document.uri, undefined, undefined, true); if (cancellation.isCancellationRequested) { return; } @@ -97,7 +97,7 @@ export class NotebookContentProvider implements INotebookContentProvider { document: NotebookDocument, cancellation: CancellationToken ): Promise { - const model = await this.notebookStorage.load(document.uri, undefined, undefined, true); + const model = await this.notebookStorage.get(document.uri, undefined, undefined, true); if (!cancellation.isCancellationRequested) { await this.notebookStorage.saveAs(model, targetResource); } @@ -107,7 +107,7 @@ export class NotebookContentProvider implements INotebookContentProvider { _context: NotebookDocumentBackupContext, cancellation: CancellationToken ): Promise { - const model = await this.notebookStorage.load(document.uri, undefined, undefined, true); + const model = await this.notebookStorage.get(document.uri, undefined, undefined, true); const id = this.notebookStorage.generateBackupId(model); await this.notebookStorage.backup(model, cancellation, id); return { diff --git a/src/client/datascience/notebook/executionService.ts b/src/client/datascience/notebook/executionService.ts index 95b2bc3fdba6..613c29a4e767 100644 --- a/src/client/datascience/notebook/executionService.ts +++ b/src/client/datascience/notebook/executionService.ts @@ -144,7 +144,7 @@ export class NotebookExecutionService implements INotebookExecutionService { private async getNotebookAndModel( document: NotebookDocument ): Promise<{ model: VSCodeNotebookModel; nb: INotebook }> { - const model = await this.notebookStorage.load(document.uri, undefined, undefined, true); + const model = await this.notebookStorage.get(document.uri, undefined, undefined, true); const nb = await this.notebookProvider.getOrCreateNotebook({ identity: document.uri, resource: document.uri, diff --git a/src/client/datascience/notebook/helpers/cellUpdateHelpers.ts b/src/client/datascience/notebook/helpers/cellUpdateHelpers.ts index c0d9274e2813..eb91df59f389 100644 --- a/src/client/datascience/notebook/helpers/cellUpdateHelpers.ts +++ b/src/client/datascience/notebook/helpers/cellUpdateHelpers.ts @@ -10,7 +10,7 @@ */ import * as assert from 'assert'; -import { NotebookCell } from '../../../../../types/vscode-proposed'; +import { NotebookCell, NotebookDocument } from '../../../../../types/vscode-proposed'; import { NotebookCellLanguageChangeEvent, NotebookCellOutputsChangeEvent, @@ -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'); @@ -51,6 +56,48 @@ export function updateCellModelWithChangesToVSCCell( } } +/** + * Updates a notebook document as a result of trusting it. + */ +export function updateVSCNotebookAfterTrustingNotebook(document: NotebookDocument, model: INotebookModel) { + const areAllCellsEditableAndRunnable = document.cells.every((cell) => { + if (cell.cellKind === vscodeNotebookEnums.CellKind.Markdown) { + return cell.metadata.editable; + } else { + return cell.metadata.editable && cell.metadata.runnable; + } + }); + const isDocumentEditableAndRunnable = + document.metadata.cellEditable && + document.metadata.cellRunnable && + document.metadata.editable && + document.metadata.runnable; + + // If already trusted, then nothing to do. + if (isDocumentEditableAndRunnable && areAllCellsEditableAndRunnable) { + return; + } + + document.metadata.cellEditable = true; + document.metadata.cellRunnable = true; + document.metadata.editable = true; + document.metadata.runnable = true; + + document.cells.forEach((cell) => { + cell.metadata.editable = true; + 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); + } + }); +} + export function clearCellForExecution(vscCell: NotebookCell, model: VSCodeNotebookModel) { vscCell.metadata.statusMessage = undefined; vscCell.metadata.executionOrder = undefined; diff --git a/src/client/datascience/notebook/helpers/helpers.ts b/src/client/datascience/notebook/helpers/helpers.ts index 99487d2f94b4..40bfc4085468 100644 --- a/src/client/datascience/notebook/helpers/helpers.ts +++ b/src/client/datascience/notebook/helpers/helpers.ts @@ -26,12 +26,12 @@ import { createCodeCell, createMarkdownCell } from '../../../../datascience-ui/c import { MARKDOWN_LANGUAGE, PYTHON_LANGUAGE } from '../../../common/constants'; import { traceError, traceWarning } from '../../../common/logger'; import { CellState, ICell, INotebookModel } from '../../types'; +import { JupyterNotebookView } from '../constants'; import { mapVSCNotebookCellToCellModel } from './cellMappers'; // tslint:disable-next-line: no-var-requires no-require-imports const vscodeNotebookEnums = require('vscode') as typeof import('vscode-proposed'); // tslint:disable-next-line: no-require-imports import cloneDeep = require('lodash/cloneDeep'); -import { JupyterNotebookView } from '../constants'; // This is the custom type we are adding into nbformat.IBaseCellMetadata interface IBaseCellVSCodeMetadata { @@ -61,11 +61,11 @@ export function notebookModelToVSCNotebookData(model: INotebookModel): NotebookD cells, languages: [defaultLanguage], metadata: { - cellEditable: true, - cellRunnable: true, - editable: true, + cellEditable: model.isTrusted, + cellRunnable: model.isTrusted, + editable: model.isTrusted, cellHasExecutionOrder: true, - runnable: true, + runnable: model.isTrusted, displayOrder: [ 'application/vnd.*', 'application/vdom.*', @@ -172,23 +172,16 @@ export function createVSCNotebookCellDataFromCell(model: INotebookModel, cell: I runState = vscodeNotebookEnums.NotebookCellRunState.Success; } - const notebookCellData: NotebookCellData = { - cellKind: - cell.data.cell_type === 'code' ? vscodeNotebookEnums.CellKind.Code : vscodeNotebookEnums.CellKind.Markdown, - language: cell.data.cell_type === 'code' ? defaultCodeLanguage : MARKDOWN_LANGUAGE, - metadata: { - editable: true, - executionOrder: typeof cell.data.execution_count === 'number' ? cell.data.execution_count : undefined, - hasExecutionOrder: cell.data.cell_type === 'code', - runState, - runnable: cell.data.cell_type === 'code' - }, - source: concatMultilineStringInput(cell.data.source), - outputs + const notebookCellMetadata: NotebookCellMetadata = { + editable: model.isTrusted, + executionOrder: typeof cell.data.execution_count === 'number' ? cell.data.execution_count : undefined, + hasExecutionOrder: cell.data.cell_type === 'code', + runState, + runnable: cell.data.cell_type === 'code' && model.isTrusted }; if (statusMessage) { - notebookCellData.metadata.statusMessage = statusMessage; + notebookCellMetadata.statusMessage = statusMessage; } const vscodeMetadata = (cell.data.metadata.vscode as unknown) as IBaseCellVSCodeMetadata | undefined; const startExecutionTime = vscodeMetadata?.start_execution_time @@ -199,12 +192,27 @@ export function createVSCNotebookCellDataFromCell(model: INotebookModel, cell: I : undefined; if (startExecutionTime && typeof endExecutionTime === 'number') { - notebookCellData.metadata.runStartTime = startExecutionTime; - notebookCellData.metadata.lastRunDuration = endExecutionTime - startExecutionTime; + notebookCellMetadata.runStartTime = startExecutionTime; + notebookCellMetadata.lastRunDuration = endExecutionTime - startExecutionTime; } - updateVSCNotebookCellMetadata(notebookCellData.metadata, cell); - return notebookCellData; + updateVSCNotebookCellMetadata(notebookCellMetadata, cell); + + // If not trusted, then clear the output in VSC Cell. + // At this point we have the original output in the ICell. + if (!model.isTrusted) { + while (outputs.length) { + outputs.shift(); + } + } + return { + cellKind: + cell.data.cell_type === 'code' ? vscodeNotebookEnums.CellKind.Code : vscodeNotebookEnums.CellKind.Markdown, + language: cell.data.cell_type === 'code' ? defaultCodeLanguage : MARKDOWN_LANGUAGE, + metadata: notebookCellMetadata, + source: concatMultilineStringInput(cell.data.source), + outputs + }; } export function createVSCCellOutputsFromOutputs(outputs?: nbformat.IOutput[]): CellOutput[] { diff --git a/src/client/datascience/notebook/notebookEditorProvider.ts b/src/client/datascience/notebook/notebookEditorProvider.ts index cdf3c0614742..bfe618d46d71 100644 --- a/src/client/datascience/notebook/notebookEditorProvider.ts +++ b/src/client/datascience/notebook/notebookEditorProvider.ts @@ -152,7 +152,7 @@ export class NotebookEditorProvider implements INotebookEditorProvider { return; } const uri = doc.uri; - const model = await this.storage.load(uri, undefined, undefined, true); + const model = await this.storage.get(uri, undefined, undefined, true); mapVSCNotebookCellsToNotebookCellModels(doc, model); // In open method we might be waiting. let editor = this.notebookEditorsByUri.get(uri.toString()); @@ -180,6 +180,9 @@ export class NotebookEditorProvider implements INotebookEditorProvider { const deferred = this.notebooksWaitingToBeOpenedByUri.get(uri.toString())!; deferred.resolve(editor); this.notebookEditorsByUri.set(uri.toString(), editor); + if (!model.isTrusted) { + await this.commandManager.executeCommand(Commands.TrustNotebook, model.file); + } } private onDidChangeActiveVsCodeNotebookEditor(editor: VSCodeNotebookEditor | undefined) { if (!editor) { @@ -229,7 +232,7 @@ export class NotebookEditorProvider implements INotebookEditorProvider { if (!isJupyterNotebook(e.document)) { return; } - const model = await this.storage.load(e.document.uri, undefined, undefined, true); + const model = await this.storage.get(e.document.uri, undefined, undefined, true); if (!(model instanceof VSCodeNotebookModel)) { throw new Error('NotebookModel not of type VSCodeNotebookModel'); } diff --git a/src/client/datascience/notebook/notebookTrustHandler.ts b/src/client/datascience/notebook/notebookTrustHandler.ts new file mode 100644 index 000000000000..06b2a864ef33 --- /dev/null +++ b/src/client/datascience/notebook/notebookTrustHandler.ts @@ -0,0 +1,38 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +'use strict'; + +import { inject, injectable } from 'inversify'; +import { IExtensionSingleActivationService } from '../../activation/types'; +import { IVSCodeNotebook } from '../../common/application/types'; +import { IFileSystem } from '../../common/platform/types'; +import { IDisposableRegistry } from '../../common/types'; +import { INotebookEditorProvider, ITrustService } from '../types'; +import { updateVSCNotebookAfterTrustingNotebook } from './helpers/cellUpdateHelpers'; +import { isJupyterNotebook } from './helpers/helpers'; + +@injectable() +export class NotebookTrustHandler implements IExtensionSingleActivationService { + constructor( + @inject(ITrustService) private readonly trustService: ITrustService, + @inject(IVSCodeNotebook) private readonly vscNotebook: IVSCodeNotebook, + @inject(INotebookEditorProvider) private readonly editorProvider: INotebookEditorProvider, + @inject(IFileSystem) private readonly fs: IFileSystem, + @inject(IDisposableRegistry) private readonly disposables: IDisposableRegistry + ) {} + public async activate(): Promise { + this.trustService.onDidSetNotebookTrust(this.onDidTrustNotebook, this, this.disposables); + } + private onDidTrustNotebook() { + this.vscNotebook.notebookDocuments.forEach((doc) => { + if (!isJupyterNotebook(doc)) { + return; + } + const editor = this.editorProvider.editors.find((e) => this.fs.arePathsSame(e.file.fsPath, doc.uri.fsPath)); + if (editor && editor.model?.isTrusted) { + updateVSCNotebookAfterTrustingNotebook(doc, editor.model); + } + }); + } +} diff --git a/src/client/datascience/notebook/serviceRegistry.ts b/src/client/datascience/notebook/serviceRegistry.ts index d8b1cea92b4f..6e53197efcb7 100644 --- a/src/client/datascience/notebook/serviceRegistry.ts +++ b/src/client/datascience/notebook/serviceRegistry.ts @@ -10,6 +10,7 @@ import { NotebookContentProvider } from './contentProvider'; import { NotebookExecutionService } from './executionService'; import { NotebookIntegration } from './integration'; import { NotebookKernel } from './notebookKernel'; +import { NotebookTrustHandler } from './notebookTrustHandler'; import { NotebookOutputRenderer } from './renderer'; import { NotebookSurveyBanner, NotebookSurveyDataLogger } from './survey'; import { INotebookContentProvider, INotebookExecutionService } from './types'; @@ -33,4 +34,8 @@ export function registerTypes(serviceManager: IServiceManager) { IExtensionSingleActivationService, NotebookSurveyDataLogger ); + serviceManager.addSingleton( + IExtensionSingleActivationService, + NotebookTrustHandler + ); } diff --git a/src/client/datascience/serviceRegistry.ts b/src/client/datascience/serviceRegistry.ts index 8fa61793f59b..31bfd9a717b3 100644 --- a/src/client/datascience/serviceRegistry.ts +++ b/src/client/datascience/serviceRegistry.ts @@ -66,6 +66,7 @@ import { NativeEditorStorage } from './interactive-ipynb/nativeEditorStorage'; import { NativeEditorSynchronizer } from './interactive-ipynb/nativeEditorSynchronizer'; import { NativeEditorViewTracker } from './interactive-ipynb/nativeEditorViewTracker'; import { INotebookStorageProvider, NotebookStorageProvider } from './interactive-ipynb/notebookStorageProvider'; +import { TrustCommandHandler } from './interactive-ipynb/trustCommandHandler'; import { TrustService } from './interactive-ipynb/trustService'; import { InteractiveWindow } from './interactive-window/interactiveWindow'; import { InteractiveWindowCommandListener } from './interactive-window/interactiveWindowCommandListener'; @@ -246,6 +247,7 @@ export function registerTypes(serviceManager: IServiceManager) { serviceManager.addSingleton(IExtensionSingleActivationService, ServerPreload); serviceManager.addSingleton(IExtensionSingleActivationService, NativeEditorViewTracker); serviceManager.addSingleton(IExtensionSingleActivationService, NotebookUsageTracker); + serviceManager.addSingleton(IExtensionSingleActivationService, TrustCommandHandler); serviceManager.addSingleton(IInteractiveWindowListener, DataScienceSurveyBannerLogger); serviceManager.addSingleton(IInteractiveWindowProvider, InteractiveWindowProvider); serviceManager.addSingleton(IJupyterDebugger, JupyterDebugger, undefined, [ICellHashListener]); diff --git a/src/client/datascience/types.ts b/src/client/datascience/types.ts index 65259318741f..d2b4142ee7f6 100644 --- a/src/client/datascience/types.ts +++ b/src/client/datascience/types.ts @@ -1043,8 +1043,8 @@ export interface INotebookStorage { save(model: INotebookModel, cancellation: CancellationToken): Promise; saveAs(model: INotebookModel, targetResource: Uri): Promise; backup(model: INotebookModel, cancellation: CancellationToken, backupId?: string): Promise; - load(file: Uri, contents?: string, backupId?: string, forVSCodeNotebook?: boolean): Promise; - load( + get(file: Uri, contents?: string, backupId?: string, forVSCodeNotebook?: boolean): Promise; + get( file: Uri, contents?: string, // tslint:disable-next-line: unified-signatures diff --git a/src/test/datascience/.vscode/settings.json b/src/test/datascience/.vscode/settings.json index 84a3ffb84a43..90f4ac4dc843 100644 --- a/src/test/datascience/.vscode/settings.json +++ b/src/test/datascience/.vscode/settings.json @@ -24,5 +24,6 @@ "files.autoSave": "off", // We don't want jupyter to start when testing (DS functionality or anything else). "python.dataScience.disableJupyterAutoStart": true, - "python.logging.level": "debug" + "python.logging.level": "debug", + "python.dataScience.alwaysTrustNotebooks": true // In UI tests we don't want prompts for `Do you want to trust thie nb...` } diff --git a/src/test/datascience/export/exportUtil.test.ts b/src/test/datascience/export/exportUtil.test.ts index 340828976e7e..a287b67ac84e 100644 --- a/src/test/datascience/export/exportUtil.test.ts +++ b/src/test/datascience/export/exportUtil.test.ts @@ -37,7 +37,7 @@ suite('DataScience - Export Util', () => { ); await exportUtil.removeSvgs(file); - const model = await notebookStorage.load(file); + const model = await notebookStorage.get(file); // make sure no svg exists in model const SVG = 'image/svg+xml'; diff --git a/src/test/datascience/interactive-common/trustCommandHandler.unit.test.ts b/src/test/datascience/interactive-common/trustCommandHandler.unit.test.ts new file mode 100644 index 000000000000..3bb98ca5c813 --- /dev/null +++ b/src/test/datascience/interactive-common/trustCommandHandler.unit.test.ts @@ -0,0 +1,127 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +'use strict'; + +import * as fakeTimers from '@sinonjs/fake-timers'; +import { assert } from 'chai'; +import * as sinon from 'sinon'; +import { anything, instance, mock, verify, when } from 'ts-mockito'; +import { Uri } from 'vscode'; +import { IExtensionSingleActivationService } from '../../../client/activation/types'; +import { IApplicationShell, ICommandManager } from '../../../client/common/application/types'; +import { ContextKey } from '../../../client/common/contextKey'; +import { EnableTrustedNotebooks } from '../../../client/common/experiments/groups'; +import { IDisposable, IExperimentService } from '../../../client/common/types'; +import { DataScience } from '../../../client/common/utils/localize'; +import { Commands } from '../../../client/datascience/constants'; +import { INotebookStorageProvider } from '../../../client/datascience/interactive-ipynb/notebookStorageProvider'; +import { TrustCommandHandler } from '../../../client/datascience/interactive-ipynb/trustCommandHandler'; +import { TrustService } from '../../../client/datascience/interactive-ipynb/trustService'; +import { INotebookEditorProvider, INotebookModel, ITrustService } from '../../../client/datascience/types'; +import { noop } from '../../core'; +import { createNotebookModel, disposeAllDisposables } from '../notebook/helper'; + +// tslint:disable: no-any + +suite('DataScience - Trust Command Handler', () => { + let trustCommandHandler: IExtensionSingleActivationService; + let trustService: ITrustService; + let editorProvider: INotebookEditorProvider; + let storageProvider: INotebookStorageProvider; + let commandManager: ICommandManager; + let applicationShell: IApplicationShell; + let disposables: IDisposable[] = []; + let clock: fakeTimers.InstalledClock; + let contextKeySet: sinon.SinonStub<[boolean], Promise>; + let experiments: IExperimentService; + let model: INotebookModel; + let trustNotebookCommandCallback: (uri: Uri) => Promise; + setup(() => { + trustService = mock(); + editorProvider = mock(); + storageProvider = mock(); + commandManager = mock(); + applicationShell = mock(); + model = createNotebookModel(false, Uri.file('a')); + when(storageProvider.get(anything())).thenResolve(model); + disposables = []; + + experiments = mock(); + + when(trustService.trustNotebook(anything(), anything())).thenResolve(); + when(experiments.inExperiment(anything())).thenCall((exp) => + Promise.resolve(exp === EnableTrustedNotebooks.experiment) + ); + when(commandManager.registerCommand(anything(), anything(), anything())).thenCall(() => ({ dispose: noop })); + when(commandManager.registerCommand(Commands.TrustNotebook, anything(), anything())).thenCall((_, cb) => { + trustNotebookCommandCallback = cb.bind(trustCommandHandler); + return { dispose: noop }; + }); + + trustCommandHandler = new TrustCommandHandler( + instance(trustService), + instance(editorProvider), + instance(storageProvider), + instance(commandManager), + instance(applicationShell), + disposables, + instance(experiments) + ); + + clock = fakeTimers.install(); + + contextKeySet = sinon.stub(ContextKey.prototype, 'set'); + contextKeySet.resolves(); + }); + teardown(() => { + sinon.restore(); + disposeAllDisposables(disposables); + clock.uninstall(); + }); + + test('Context not set if not in experiment', async () => { + when(experiments.inExperiment(anything())).thenResolve(false); + + await trustCommandHandler.activate(); + await clock.runAllAsync(); + + assert.equal(contextKeySet.callCount, 0); + }); + test('Context set if in experiment', async () => { + when(experiments.inExperiment(anything())).thenCall((exp) => + Promise.resolve(exp === EnableTrustedNotebooks.experiment) + ); + + await trustCommandHandler.activate(); + await clock.runAllAsync(); + + assert.equal(contextKeySet.callCount, 1); + }); + test('Executing command will not update trust after dismissing the prompt', async () => { + when(applicationShell.showErrorMessage(anything(), anything(), anything(), anything())).thenResolve( + undefined as any + ); + + await trustCommandHandler.activate(); + await clock.runAllAsync(); + await trustNotebookCommandCallback(Uri.file('a')); + + verify(applicationShell.showErrorMessage(anything(), anything(), anything(), anything())).once(); + verify(trustService.trustNotebook(anything(), anything())).never(); + assert.isFalse(model.isTrusted); + }); + test('Executing command will update trust', async () => { + when(applicationShell.showErrorMessage(anything(), anything(), anything(), anything())).thenResolve( + DataScience.trustNotebook() as any + ); + + assert.isFalse(model.isTrusted); + await trustCommandHandler.activate(); + await clock.runAllAsync(); + await trustNotebookCommandCallback(Uri.file('a')); + + verify(applicationShell.showErrorMessage(anything(), anything(), anything(), anything())).once(); + verify(trustService.trustNotebook(anything(), anything())).once(); + assert.isTrue(model.isTrusted); + }); +}); diff --git a/src/test/datascience/trustService.unit.test.ts b/src/test/datascience/interactive-common/trustService.unit.test.ts similarity index 84% rename from src/test/datascience/trustService.unit.test.ts rename to src/test/datascience/interactive-common/trustService.unit.test.ts index 2ff31348abf6..beb3c7c88fad 100644 --- a/src/test/datascience/trustService.unit.test.ts +++ b/src/test/datascience/interactive-common/trustService.unit.test.ts @@ -8,12 +8,12 @@ import * as os from 'os'; import { anything, instance, mock, when } from 'ts-mockito'; import * as typemoq from 'typemoq'; import { Uri } from 'vscode'; -import { ConfigurationService } from '../../client/common/configuration/service'; -import { ExperimentService } from '../../client/common/experiments/service'; -import { FileSystem } from '../../client/common/platform/fileSystem'; -import { IExtensionContext } from '../../client/common/types'; -import { DigestStorage } from '../../client/datascience/interactive-ipynb/digestStorage'; -import { TrustService } from '../../client/datascience/interactive-ipynb/trustService'; +import { ConfigurationService } from '../../../client/common/configuration/service'; +import { ExperimentService } from '../../../client/common/experiments/service'; +import { FileSystem } from '../../../client/common/platform/fileSystem'; +import { IExtensionContext } from '../../../client/common/types'; +import { DigestStorage } from '../../../client/datascience/interactive-ipynb/digestStorage'; +import { TrustService } from '../../../client/datascience/interactive-ipynb/trustService'; suite('DataScience - TrustService', () => { let trustService: TrustService; diff --git a/src/test/datascience/interactive-ipynb/nativeEditorStorage.unit.test.ts b/src/test/datascience/interactive-ipynb/nativeEditorStorage.unit.test.ts index 98b7a9fae312..3db2eecd0ee4 100644 --- a/src/test/datascience/interactive-ipynb/nativeEditorStorage.unit.test.ts +++ b/src/test/datascience/interactive-ipynb/nativeEditorStorage.unit.test.ts @@ -409,7 +409,7 @@ suite('DataScience - Native Editor Storage', () => { } test('Create new editor and add some cells', async () => { - model = await storage.load(baseUri); + model = await storage.get(baseUri); insertCell(0, '1'); const cells = model.cells; expect(cells).to.be.lengthOf(4); @@ -418,7 +418,7 @@ suite('DataScience - Native Editor Storage', () => { }); test('Move cells around', async () => { - model = await storage.load(baseUri); + model = await storage.get(baseUri); swapCells('NotebookImport#0', 'NotebookImport#1'); const cells = model.cells; expect(cells).to.be.lengthOf(3); @@ -427,7 +427,7 @@ suite('DataScience - Native Editor Storage', () => { }); test('Edit/delete cells', async () => { - model = await storage.load(baseUri); + model = await storage.get(baseUri); expect(model.isDirty).to.be.equal(false, 'Editor should not be dirty'); editCell( [ @@ -467,7 +467,7 @@ suite('DataScience - Native Editor Storage', () => { test('Editing a file and closing will keep contents', async () => { await filesConfig?.update('autoSave', 'off'); - model = await storage.load(baseUri); + model = await storage.get(baseUri); expect(model.isDirty).to.be.equal(false, 'Editor should not be dirty'); editCell( [ @@ -496,7 +496,7 @@ suite('DataScience - Native Editor Storage', () => { // Recreate storage = createStorage(); - model = await storage.load(baseUri); + model = await storage.get(baseUri); const cells = model.cells; expect(cells).to.be.lengthOf(3); @@ -506,7 +506,7 @@ suite('DataScience - Native Editor Storage', () => { }); test('Editing a new file and closing will keep contents', async () => { - model = await storage.load(untiledUri, undefined, true); + model = await storage.get(untiledUri, undefined, true); expect(model.isDirty).to.be.equal(false, 'Editor should not be dirty'); insertCell(0, 'a=1'); @@ -515,7 +515,7 @@ suite('DataScience - Native Editor Storage', () => { // Recreate storage = createStorage(); - model = await storage.load(untiledUri); + model = await storage.get(untiledUri); const cells = model.cells; expect(cells).to.be.lengthOf(2); @@ -534,7 +534,7 @@ suite('DataScience - Native Editor Storage', () => { // Put the regular file into the local storage await localMemento.update(`notebook-storage-${file.toString()}`, differentFile); - model = await storage.load(file); + model = await storage.get(file); // It should load with that value const cells = model.cells; @@ -555,7 +555,7 @@ suite('DataScience - Native Editor Storage', () => { contents: differentFile, lastModifiedTimeMs: Date.now() }); - model = await storage.load(file); + model = await storage.get(file); // It should load with that value const cells = model.cells; @@ -585,7 +585,7 @@ suite('DataScience - Native Editor Storage', () => { lastModifiedTimeMs: Date.now() }); - model = await storage.load(file); + model = await storage.get(file); // It should load with that value const cells = model.cells; diff --git a/src/test/datascience/notebook/contentProvider.unit.test.ts b/src/test/datascience/notebook/contentProvider.unit.test.ts index b7c27604c456..b85a4715ebcf 100644 --- a/src/test/datascience/notebook/contentProvider.unit.test.ts +++ b/src/test/datascience/notebook/contentProvider.unit.test.ts @@ -33,157 +33,173 @@ suite('Data Science - NativeNotebook ContentProvider', () => { instance(compatSupport) ); }); + [true, false].forEach((isNotebookTrusted) => { + suite(isNotebookTrusted ? 'Trusted Notebook' : 'Un-trusted notebook', () => { + test('Return notebook with 2 cells', async () => { + const model: Partial = { + cells: [ + { + data: { + cell_type: 'code', + execution_count: 10, + hasExecutionOrder: true, + outputs: [], + source: 'print(1)', + metadata: {} + }, + file: 'a.ipynb', + id: 'MyCellId1', + line: 0, + state: CellState.init + }, + { + data: { + cell_type: 'markdown', + hasExecutionOrder: false, + source: '# HEAD', + metadata: {} + }, + file: 'a.ipynb', + id: 'MyCellId2', + line: 0, + state: CellState.init + } + ], + isTrusted: isNotebookTrusted + }; + when(storageProvider.get(anything(), anything(), anything(), anything())).thenResolve( + (model as unknown) as INotebookModel + ); - test('Return notebook with 2 cells', async () => { - const model: Partial = { - cells: [ - { - data: { - cell_type: 'code', - execution_count: 10, - hasExecutionOrder: true, + const notebook = await contentProvider.openNotebook(fileUri, {}); + + assert.isOk(notebook); + assert.deepEqual(notebook.languages, [PYTHON_LANGUAGE]); + // ignore metadata we add. + notebook.cells.forEach((cell) => delete cell.metadata.custom); + + assert.equal(notebook.metadata.cellEditable, isNotebookTrusted); + assert.equal(notebook.metadata.cellRunnable, isNotebookTrusted); + assert.equal(notebook.metadata.editable, isNotebookTrusted); + assert.equal(notebook.metadata.runnable, isNotebookTrusted); + + assert.deepEqual(notebook.cells, [ + { + cellKind: (vscodeNotebookEnums as any).CellKind.Code, + language: PYTHON_LANGUAGE, outputs: [], source: 'print(1)', - metadata: {} + metadata: { + editable: isNotebookTrusted, + executionOrder: 10, + hasExecutionOrder: true, + runState: (vscodeNotebookEnums as any).NotebookCellRunState.Success, + runnable: isNotebookTrusted + } }, - file: 'a.ipynb', - id: 'MyCellId1', - line: 0, - state: CellState.init - }, - { - data: { - cell_type: 'markdown', - hasExecutionOrder: false, + { + cellKind: (vscodeNotebookEnums as any).CellKind.Markdown, + language: MARKDOWN_LANGUAGE, + outputs: [], source: '# HEAD', - metadata: {} + metadata: { + editable: isNotebookTrusted, + executionOrder: undefined, + hasExecutionOrder: false, + runState: (vscodeNotebookEnums as any).NotebookCellRunState.Idle, + runnable: false + } + } + ]); + }); + + test('Return notebook with csharp language', async () => { + const model: Partial = { + metadata: { + language_info: { + name: 'csharp' + }, + orig_nbformat: 5 }, - file: 'a.ipynb', - id: 'MyCellId2', - line: 0, - state: CellState.init - } - ] - }; - when(storageProvider.load(anything(), anything(), anything(), anything())).thenResolve( - (model as unknown) as INotebookModel - ); + cells: [ + { + data: { + cell_type: 'code', + execution_count: 10, + hasExecutionOrder: true, + outputs: [], + source: 'Console.WriteLine("1")', + metadata: {} + }, + file: 'a.ipynb', + id: 'MyCellId1', + line: 0, + state: CellState.init + }, + { + data: { + cell_type: 'markdown', + hasExecutionOrder: false, + source: '# HEAD', + metadata: {} + }, + file: 'a.ipynb', + id: 'MyCellId2', + line: 0, + state: CellState.init + } + ], + isTrusted: isNotebookTrusted + }; + when(storageProvider.get(anything(), anything(), anything(), anything())).thenResolve( + (model as unknown) as INotebookModel + ); - const notebook = await contentProvider.openNotebook(fileUri, {}); + const notebook = await contentProvider.openNotebook(fileUri, {}); - assert.isOk(notebook); - assert.deepEqual(notebook.languages, [PYTHON_LANGUAGE]); - // ignore metadata we add. - notebook.cells.forEach((cell) => delete cell.metadata.custom); + assert.isOk(notebook); + assert.deepEqual(notebook.languages, ['csharp']); - assert.deepEqual(notebook.cells, [ - { - cellKind: (vscodeNotebookEnums as any).CellKind.Code, - language: PYTHON_LANGUAGE, - outputs: [], - source: 'print(1)', - metadata: { - editable: true, - executionOrder: 10, - hasExecutionOrder: true, - runState: (vscodeNotebookEnums as any).NotebookCellRunState.Success, - runnable: true - } - }, - { - cellKind: (vscodeNotebookEnums as any).CellKind.Markdown, - language: MARKDOWN_LANGUAGE, - outputs: [], - source: '# HEAD', - metadata: { - editable: true, - executionOrder: undefined, - hasExecutionOrder: false, - runState: (vscodeNotebookEnums as any).NotebookCellRunState.Idle, - runnable: false - } - } - ]); - }); + assert.equal(notebook.metadata.cellEditable, isNotebookTrusted); + assert.equal(notebook.metadata.cellRunnable, isNotebookTrusted); + assert.equal(notebook.metadata.editable, isNotebookTrusted); + assert.equal(notebook.metadata.runnable, isNotebookTrusted); + + // ignore metadata we add. + notebook.cells.forEach((cell: NotebookCellData) => delete cell.metadata.custom); - test('Return notebook with csharp language', async () => { - const model: Partial = { - metadata: { - language_info: { - name: 'csharp' - }, - orig_nbformat: 5 - }, - cells: [ - { - data: { - cell_type: 'code', - execution_count: 10, - hasExecutionOrder: true, + assert.deepEqual(notebook.cells, [ + { + cellKind: (vscodeNotebookEnums as any).CellKind.Code, + language: 'csharp', outputs: [], source: 'Console.WriteLine("1")', - metadata: {} + metadata: { + editable: isNotebookTrusted, + executionOrder: 10, + hasExecutionOrder: true, + runState: (vscodeNotebookEnums as any).NotebookCellRunState.Success, + runnable: isNotebookTrusted + } }, - file: 'a.ipynb', - id: 'MyCellId1', - line: 0, - state: CellState.init - }, - { - data: { - cell_type: 'markdown', - hasExecutionOrder: false, + { + cellKind: (vscodeNotebookEnums as any).CellKind.Markdown, + language: MARKDOWN_LANGUAGE, + outputs: [], source: '# HEAD', - metadata: {} - }, - file: 'a.ipynb', - id: 'MyCellId2', - line: 0, - state: CellState.init - } - ] - }; - when(storageProvider.load(anything(), anything(), anything(), anything())).thenResolve( - (model as unknown) as INotebookModel - ); - - const notebook = await contentProvider.openNotebook(fileUri, {}); - - assert.isOk(notebook); - assert.deepEqual(notebook.languages, ['csharp']); - // ignore metadata we add. - notebook.cells.forEach((cell: NotebookCellData) => delete cell.metadata.custom); - - assert.deepEqual(notebook.cells, [ - { - cellKind: (vscodeNotebookEnums as any).CellKind.Code, - language: 'csharp', - outputs: [], - source: 'Console.WriteLine("1")', - metadata: { - editable: true, - executionOrder: 10, - hasExecutionOrder: true, - runState: (vscodeNotebookEnums as any).NotebookCellRunState.Success, - runnable: true - } - }, - { - cellKind: (vscodeNotebookEnums as any).CellKind.Markdown, - language: MARKDOWN_LANGUAGE, - outputs: [], - source: '# HEAD', - metadata: { - editable: true, - executionOrder: undefined, - hasExecutionOrder: false, - runState: (vscodeNotebookEnums as any).NotebookCellRunState.Idle, - runnable: false - } - } - ]); - }); - test('Verify mime types and order', () => { - // https://github.com/microsoft/vscode-python/issues/11880 + metadata: { + editable: isNotebookTrusted, + executionOrder: undefined, + hasExecutionOrder: false, + runState: (vscodeNotebookEnums as any).NotebookCellRunState.Idle, + runnable: false + } + } + ]); + }); + test('Verify mime types and order', () => { + // https://github.com/microsoft/vscode-python/issues/11880 + }); + }); }); }); diff --git a/src/test/datascience/notebook/helper.ts b/src/test/datascience/notebook/helper.ts index 8f1d66148b59..cb77fdb87749 100644 --- a/src/test/datascience/notebook/helper.ts +++ b/src/test/datascience/notebook/helper.ts @@ -9,16 +9,30 @@ import * as fs from 'fs-extra'; import * as path from 'path'; import * as sinon from 'sinon'; import * as tmp from 'tmp'; -import { commands, Uri } from 'vscode'; -import { NotebookCell } from '../../../../types/vscode-proposed'; +import { instance, mock } from 'ts-mockito'; +import { commands, TextDocument, Uri } from 'vscode'; +import { NotebookCell, NotebookDocument } from '../../../../types/vscode-proposed'; import { CellDisplayOutput } from '../../../../typings/vscode-proposed'; import { IApplicationEnvironment, IVSCodeNotebook } from '../../../client/common/application/types'; import { MARKDOWN_LANGUAGE, PYTHON_LANGUAGE } from '../../../client/common/constants'; import { IDisposable } from '../../../client/common/types'; import { noop, swallowExceptions } from '../../../client/common/utils/misc'; -import { findMappedNotebookCellModel } from '../../../client/datascience/notebook/helpers/cellMappers'; +import { Identifiers } from '../../../client/datascience/constants'; +import { JupyterNotebookView } from '../../../client/datascience/notebook/constants'; +import { + findMappedNotebookCellModel, + mapVSCNotebookCellsToNotebookCellModels +} from '../../../client/datascience/notebook/helpers/cellMappers'; +import { createVSCNotebookCellDataFromCell } from '../../../client/datascience/notebook/helpers/helpers'; import { INotebookContentProvider } from '../../../client/datascience/notebook/types'; -import { ICell, INotebookEditorProvider, INotebookModel, INotebookProvider } from '../../../client/datascience/types'; +import { VSCodeNotebookModel } from '../../../client/datascience/notebookStorage/vscNotebookModel'; +import { + CellState, + ICell, + INotebookEditorProvider, + INotebookModel, + INotebookProvider +} from '../../../client/datascience/types'; import { createEventHandler, waitForCondition } from '../../common'; import { EXTENSION_ROOT_DIR_FOR_TESTS } from '../../constants'; import { closeActiveWindows, initialize } from '../../initialize'; @@ -365,3 +379,61 @@ export async function saveActiveNotebook(disposables: IDisposable[]) { await waitForCondition(async () => savedEvent.all.some((e) => e.kind === 'save'), 5_000, 'Not saved'); } + +export function createNotebookModel(trusted: boolean, uri: Uri, nb?: Partial) { + const nbJson: nbformat.INotebookContent = { + cells: [], + metadata: { + orig_nbformat: 4 + }, + nbformat: 4, + nbformat_minor: 4, + ...(nb || {}) + }; + + const cells = nbJson.cells.map((c, index) => { + return { + id: `NotebookImport#${index}`, + file: Identifiers.EmptyFileName, + line: 0, + state: CellState.finished, + data: c + }; + }); + return new VSCodeNotebookModel(trusted, uri, JSON.parse(JSON.stringify(cells))); +} + +export function createNotebookDocument( + model: INotebookModel, + viewType: string = JupyterNotebookView +): NotebookDocument { + const doc: NotebookDocument = { + cells: [], + fileName: model.file.fsPath, + isDirty: false, + languages: [], + uri: model.file, + viewType, + metadata: { + cellEditable: model.isTrusted, + cellHasExecutionOrder: true, + cellRunnable: model.isTrusted, + editable: model.isTrusted, + runnable: model.isTrusted + } + }; + model.cells.forEach((cell, index) => { + const vscCell = createVSCNotebookCellDataFromCell(model, cell)!; + const vscDocumentCell: NotebookCell = { + ...vscCell, + uri: model.file.with({ fragment: `cell${index}` }), + notebook: doc, + document: instance(mock()) + }; + doc.cells.push(vscDocumentCell); + }); + if (viewType === JupyterNotebookView) { + mapVSCNotebookCellsToNotebookCellModels(doc, model); + } + return doc; +} diff --git a/src/test/datascience/notebook/helpers.unit.test.ts b/src/test/datascience/notebook/helpers.unit.test.ts index 63174224c3b3..bb48bc6c30c2 100644 --- a/src/test/datascience/notebook/helpers.unit.test.ts +++ b/src/test/datascience/notebook/helpers.unit.test.ts @@ -40,7 +40,8 @@ suite('Data Science - NativeNotebook helpers', () => { line: 0, state: CellState.init } - ] + ], + isTrusted: true }; const notebook = notebookModelToVSCNotebookData((model as unknown) as INotebookModel); @@ -95,7 +96,8 @@ suite('Data Science - NativeNotebook helpers', () => { line: 0, state: CellState.init } - ] + ], + isTrusted: true }; const notebook = notebookModelToVSCNotebookData((model as unknown) as INotebookModel); diff --git a/src/test/datascience/notebook/notebookTrustHandler.unit.test.ts b/src/test/datascience/notebook/notebookTrustHandler.unit.test.ts new file mode 100644 index 000000000000..7d3f6e821b4e --- /dev/null +++ b/src/test/datascience/notebook/notebookTrustHandler.unit.test.ts @@ -0,0 +1,142 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +'use strict'; + +import { nbformat } from '@jupyterlab/coreutils'; +import { assert } from 'chai'; +import { teardown } from 'mocha'; +import { anything, instance, mock, when } from 'ts-mockito'; +import { EventEmitter, Uri } from 'vscode'; +import { NotebookDocument } from '../../../../types/vscode-proposed'; +import { IExtensionSingleActivationService } from '../../../client/activation/types'; +import { IVSCodeNotebook } from '../../../client/common/application/types'; +import { IFileSystem } from '../../../client/common/platform/types'; +import { IDisposable } from '../../../client/common/types'; +import { NotebookTrustHandler } from '../../../client/datascience/notebook/notebookTrustHandler'; +import { INotebookEditor, INotebookEditorProvider, ITrustService } from '../../../client/datascience/types'; +import { createNotebookDocument, createNotebookModel, disposeAllDisposables } from './helper'; +// tslint:disable-next-line: no-var-requires no-require-imports +const vscodeNotebookEnums = require('vscode') as typeof import('vscode-proposed'); + +// tslint:disable: no-any +suite('Data Science - NativeNotebook TrustHandler', () => { + let trustHandler: IExtensionSingleActivationService; + let trustService: ITrustService; + let vscNotebook: IVSCodeNotebook; + let editorProvider: INotebookEditorProvider; + let fs: IFileSystem; + let disposables: IDisposable[]; + let onDidTrustNotebook: EventEmitter; + setup(async () => { + disposables = []; + trustService = mock(); + vscNotebook = mock(); + editorProvider = mock(); + fs = mock(); + onDidTrustNotebook = new EventEmitter(); + when(trustService.onDidSetNotebookTrust).thenReturn(onDidTrustNotebook.event); + when(fs.arePathsSame(anything(), anything())).thenCall((a, b) => a === b); // Dirty simple file compare. + trustHandler = new NotebookTrustHandler( + instance(trustService), + instance(vscNotebook), + instance(editorProvider), + instance(fs), + disposables + ); + + await trustHandler.activate(); + }); + teardown(() => disposeAllDisposables(disposables)); + function assertDocumentTrust(document: NotebookDocument, trusted: boolean) { + assert.equal(document.metadata.cellEditable, trusted); + assert.equal(document.metadata.cellRunnable, trusted); + assert.equal(document.metadata.editable, trusted); + assert.equal(document.metadata.runnable, trusted); + + document.cells.forEach((cell) => { + assert.equal(cell.metadata.editable, trusted); + if (cell.cellKind === vscodeNotebookEnums.CellKind.Code) { + assert.equal(cell.metadata.runnable, trusted); + // In our test all code cells have outputs. + if (trusted) { + assert.ok(cell.outputs.length, 'No output in trusted cell'); + } else { + assert.lengthOf(cell.outputs, 0, 'Cannot have output in non-trusted notebook'); + } + } + }); + } + function createModels() { + const nbJson: Partial = { + cells: [ + { + cell_type: 'markdown', + source: [], + metadata: {} + }, + { + cell_type: 'code', + source: [], + metadata: {}, + execution_count: 1, + outputs: [ + { + output_type: 'stream', + name: 'stdout', + text: 'Hello World' + } + ] + } + ] + }; + + return [createNotebookModel(false, Uri.file('a'), nbJson), createNotebookModel(false, Uri.file('b'), nbJson)]; + } + test('When a notebook is trusted, the Notebook document is updated accordingly', async () => { + const [model1, model2] = createModels(); + const [nb1, nb2, nbAnotherExtension] = [ + createNotebookDocument(model1), + createNotebookDocument(model2), + createNotebookDocument(model2, 'AnotherExtensionNotebookEditorForIpynbFile') + ]; + + // Initially un-trusted. + assertDocumentTrust(nb1, false); + assertDocumentTrust(nb2, false); + assertDocumentTrust(nbAnotherExtension, false); + + when(vscNotebook.notebookDocuments).thenReturn([nb1, nb2]); + const editor1 = mock(); + const editor2 = mock(); + when(editor1.file).thenReturn(model1.file); + when(editor2.file).thenReturn(model2.file); + when(editor1.model).thenReturn(model1); + when(editor2.model).thenReturn(model2); + when(editorProvider.editors).thenReturn([instance(editor1), instance(editor2)]); + + // Trigger a change, even though none of the models are still trusted. + onDidTrustNotebook.fire(); + + // Still un-trusted. + assertDocumentTrust(nb1, false); + assertDocumentTrust(nb2, false); + assertDocumentTrust(nbAnotherExtension, false); + + // Trigger a change, after trusting second nb/model. + model2.update({ + source: 'user', + kind: 'updateTrust', + oldDirty: model2.isDirty, + newDirty: model2.isDirty, + isNotebookTrusted: true + }); + + onDidTrustNotebook.fire(); + + // Nb1 is still un-trusted and nb1 is trusted. + assertDocumentTrust(nb1, false); + assertDocumentTrust(nb2, true); + assertDocumentTrust(nbAnotherExtension, false); // This is a document from a different content provider, we should modify this. + }); +});