Skip to content

Commit 352f2c5

Browse files
authored
Support Formatting notebooks (#12199)
For #12195 When formatting notebook cells we must create temporary files Added tests for formatting new and existing notebooks
1 parent 4efda9d commit 352f2c5

7 files changed

Lines changed: 158 additions & 11 deletions

File tree

news/1 Enhancements/12195.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Support formatting of Notebook Cells when using the VS Code Insiders API for Notebooks.

src/client/common/editor.ts

Lines changed: 13 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,10 @@ import { EOL } from 'os';
55
import * as path from 'path';
66
import { Position, Range, TextDocument, TextEdit, Uri, WorkspaceEdit } from 'vscode';
77
import { IFileSystem } from '../common/platform/types';
8+
import { traceError } from '../logging';
9+
import { WrappedError } from './errors/errorUtils';
810
import { IEditorUtils } from './types';
11+
import { isNotebookCell } from './utils/misc';
912

1013
// Code borrowed from goFormat.ts (Go Extension for VS Code)
1114
enum EditAction {
@@ -245,19 +248,25 @@ function getTextEditsInternal(before: string, diffs: [number, string][], startLi
245248
}
246249

247250
export async function getTempFileWithDocumentContents(document: TextDocument, fs: IFileSystem): Promise<string> {
248-
const ext = path.extname(document.uri.fsPath);
251+
// Hardcode extension to `.py` for notebook cells (do not use `.ipyn` doesn't work).
252+
const ext = isNotebookCell(document.uri) ? '.py' : path.extname(document.uri.fsPath);
249253
// Don't create file in temp folder since external utilities
250254
// look into configuration files in the workspace and are not
251255
// to find custom rules if file is saved in a random disk location.
252256
// This means temp file has to be created in the same folder
253257
// as the original one and then removed.
254258

255-
// tslint:disable-next-line:no-require-imports
256-
const fileName = `${document.uri.fsPath}.${md5(document.uri.fsPath)}${ext}`;
259+
let fileName = `${document.uri.fsPath}.${md5(document.uri.fsPath)}${ext}`;
260+
257261
try {
262+
// When dealing with untitled notebooks, there's no original physical file, hence create a temp file.
263+
if (isNotebookCell(document.uri) && !(await fs.fileExists(document.uri.fsPath))) {
264+
fileName = (await fs.createTemporaryFile(`${path.basename(document.uri.fsPath)}${ext}`)).filePath;
265+
}
258266
await fs.writeFile(fileName, document.getText());
259267
} catch (ex) {
260-
throw Error(`Failed to create a temporary file, ${ex.message}`);
268+
traceError('Failed to create a temporary file', ex);
269+
throw new WrappedError(`Failed to create a temporary file, ${ex.message}`, ex);
261270
}
262271
return fileName;
263272
}

src/client/common/utils/misc.ts

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -78,3 +78,7 @@ export function isNotebookCell(documentOrUri: TextDocument | Uri): boolean {
7878
const uri = isUri(documentOrUri) ? documentOrUri : documentOrUri.uri;
7979
return uri.scheme.includes(NotebookCellScheme);
8080
}
81+
82+
export function isUntitledFile(file?: Uri) {
83+
return file?.scheme === 'untitled';
84+
}

src/client/datascience/interactive-common/interactiveBase.ts

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,7 @@ import { IFileSystem } from '../../common/platform/types';
4141
import { IConfigurationService, IDisposableRegistry, IExperimentsManager } from '../../common/types';
4242
import { createDeferred, Deferred } from '../../common/utils/async';
4343
import * as localize from '../../common/utils/localize';
44-
import { noop } from '../../common/utils/misc';
44+
import { isUntitledFile, noop } from '../../common/utils/misc';
4545
import { StopWatch } from '../../common/utils/stopWatch';
4646
import { PythonInterpreter } from '../../pythonEnvironments/discovery/types';
4747
import { captureTelemetry, sendTelemetryEvent } from '../../telemetry';
@@ -65,7 +65,6 @@ import {
6565
SysInfoReason,
6666
VariableExplorerStateKeys
6767
} from '../interactive-common/interactiveWindowTypes';
68-
import { isUntitledFile } from '../interactive-ipynb/nativeEditorStorage';
6968
import { JupyterInvalidKernelError } from '../jupyter/jupyterInvalidKernelError';
7069
import { JupyterKernelPromiseFailedError } from '../jupyter/kernels/jupyterKernelPromiseFailedError';
7170
import { KernelSpecInterpreter } from '../jupyter/kernels/kernelSelector';

src/client/datascience/interactive-ipynb/nativeEditorStorage.ts

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ import { createCodeCell } from '../../../datascience-ui/common/cellFactory';
99
import { traceError } from '../../common/logger';
1010
import { IFileSystem } from '../../common/platform/types';
1111
import { GLOBAL_MEMENTO, ICryptoUtils, IExtensionContext, IMemento, WORKSPACE_MEMENTO } from '../../common/types';
12-
import { noop } from '../../common/utils/misc';
12+
import { isUntitledFile, noop } from '../../common/utils/misc';
1313
import { PythonInterpreter } from '../../pythonEnvironments/discovery/types';
1414
import { Identifiers, KnownNotebookLanguages, Telemetry } from '../constants';
1515
import { IEditorContentChange, NotebookModelChange } from '../interactive-common/interactiveWindowTypes';
@@ -35,9 +35,6 @@ interface INativeEditorStorageState {
3535
notebookJson: Partial<nbformat.INotebookContent>;
3636
}
3737

38-
export function isUntitledFile(file?: Uri) {
39-
return file?.scheme === 'untitled';
40-
}
4138
export function isUntitled(model?: INotebookModel): boolean {
4239
return isUntitledFile(model?.file);
4340
}

src/client/formatters/baseFormatter.ts

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import { traceError } from '../common/logger';
88
import { IFileSystem } from '../common/platform/types';
99
import { IPythonToolExecutionService } from '../common/process/types';
1010
import { IDisposableRegistry, IInstaller, IOutputChannel, Product } from '../common/types';
11+
import { isNotebookCell } from '../common/utils/misc';
1112
import { IServiceContainer } from '../ioc/types';
1213
import { getTempFileWithDocumentContents, getTextEditsFromPatch } from './../common/editor';
1314
import { IFormatterHelper } from './types';
@@ -61,6 +62,7 @@ export abstract class BaseFormatter {
6162
// However they don't support returning the diff of the formatted text when reading data from the input stream.
6263
// Yet getting text formatted that way avoids having to create a temporary file, however the diffing will have
6364
// to be done here in node (extension), i.e. extension CPU, i.e. less responsive solution.
65+
// Also, always create temp files for Notebook cells.
6466
const tempFile = await this.createTempFile(document);
6567
if (this.checkCancellation(document.fileName, tempFile, token)) {
6668
return [];
@@ -117,9 +119,15 @@ export abstract class BaseFormatter {
117119
this.outputChannel.appendLine(`\n${customError}\n${error}`);
118120
}
119121

122+
/**
123+
* Always create a temporary file when formatting notebook cells.
124+
* This is because there is no physical file associated with notebook cells (they are all virtual).
125+
*/
120126
private async createTempFile(document: vscode.TextDocument): Promise<string> {
121127
const fs = this.serviceContainer.get<IFileSystem>(IFileSystem);
122-
return document.isDirty ? getTempFileWithDocumentContents(document, fs) : document.fileName;
128+
return document.isDirty || isNotebookCell(document)
129+
? getTempFileWithDocumentContents(document, fs)
130+
: document.fileName;
123131
}
124132

125133
private deleteTempFile(originalFile: string, tempFile: string): Promise<void> {
Lines changed: 129 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,129 @@
1+
// Copyright (c) Microsoft Corporation. All rights reserved.
2+
// Licensed under the MIT License.
3+
4+
'use strict';
5+
6+
import { assert } from 'chai';
7+
import * as path from 'path';
8+
import * as sinon from 'sinon';
9+
import { commands, ConfigurationTarget, Uri } from 'vscode';
10+
import { IWorkspaceService } from '../../client/common/application/types';
11+
import { IFileSystem } from '../../client/common/platform/types';
12+
import { IPythonToolExecutionService } from '../../client/common/process/types';
13+
import { IDisposable, Product } from '../../client/common/types';
14+
import { INotebookEditorProvider } from '../../client/datascience/types';
15+
import { IExtensionTestApi } from '../common';
16+
import {
17+
canRunTests,
18+
closeNotebooksAndCleanUpAfterTests,
19+
createTemporaryNotebook,
20+
disposeAllDisposables,
21+
swallowSavingOfNotebooks
22+
} from '../datascience/notebook/helper';
23+
import { EXTENSION_ROOT_DIR_FOR_TESTS, initialize, initializeTest } from '../initialize';
24+
25+
// tslint:disable: no-any no-invalid-this
26+
suite('Formatting - Notebooks', () => {
27+
let api: IExtensionTestApi;
28+
suiteSetup(async function () {
29+
api = await initialize();
30+
if (!(await canRunTests())) {
31+
return this.skip();
32+
}
33+
});
34+
suiteTeardown(closeNotebooksAndCleanUpAfterTests);
35+
['yapf', 'black', 'autopep8'].forEach((formatter) => {
36+
suite(formatter, () => {
37+
const disposables: IDisposable[] = [];
38+
let testIPynb: Uri;
39+
let executionService: IPythonToolExecutionService;
40+
let editorProvider: INotebookEditorProvider;
41+
let fs: IFileSystem;
42+
const product: Product =
43+
formatter === 'yapf' ? Product.yapf : formatter === 'black' ? Product.black : Product.autopep8;
44+
suiteSetup(async () => {
45+
const workspaceService = api.serviceContainer.get<IWorkspaceService>(IWorkspaceService);
46+
const config = workspaceService.getConfiguration(
47+
'python.formatting',
48+
workspaceService.workspaceFolders![0].uri
49+
);
50+
await config.update('provider', formatter, ConfigurationTarget.Workspace);
51+
});
52+
setup(async () => {
53+
sinon.restore();
54+
await initializeTest();
55+
await swallowSavingOfNotebooks();
56+
57+
// Don't use same file (due to dirty handling, we might save in dirty.)
58+
// Cuz we won't save to file, hence extension will backup in dirty file and when u re-open it will open from dirty.
59+
const templateIPynb = path.join(
60+
EXTENSION_ROOT_DIR_FOR_TESTS,
61+
'src',
62+
'test',
63+
'datascience',
64+
'test.ipynb'
65+
);
66+
testIPynb = Uri.file(await createTemporaryNotebook(templateIPynb, disposables));
67+
executionService = api.serviceContainer.get<IPythonToolExecutionService>(IPythonToolExecutionService);
68+
editorProvider = api.serviceContainer.get<INotebookEditorProvider>(INotebookEditorProvider);
69+
fs = api.serviceContainer.get<IFileSystem>(IFileSystem);
70+
});
71+
teardown(closeNotebooksAndCleanUpAfterTests);
72+
suiteTeardown(() => disposeAllDisposables(disposables));
73+
test('Formatted with temporary file when formatting existing saved notebooks (without changes)', async () => {
74+
// Open a new notebook & add a cell
75+
await editorProvider.open(testIPynb);
76+
77+
// Check if a temp file is created.
78+
const spiedExecution = sinon.spy(executionService, 'exec');
79+
const spiedWriteFile = sinon.spy(fs, 'writeFile');
80+
disposables.push({ dispose: () => spiedWriteFile.restore() });
81+
disposables.push({ dispose: () => spiedExecution.restore() });
82+
83+
// Format the cell
84+
await commands.executeCommand('notebook.formatCell');
85+
86+
// Verify a temp file was created, having a file starting with thenb file name.
87+
assert.isOk(
88+
spiedWriteFile
89+
.getCalls()
90+
.some(
91+
(call) =>
92+
call.args[0].includes(path.basename(testIPynb.fsPath)) &&
93+
call.args[0].includes('.ipynb')
94+
),
95+
'Temp file not created'
96+
);
97+
// Verify we tried to format.
98+
assert.isOk(
99+
spiedExecution.getCalls().some((call) => call.args[0].product === product),
100+
'Not formatted'
101+
);
102+
});
103+
test('Formatted with temporary file when formatting untitled notebooks', async () => {
104+
// Open a new notebook & add a cell
105+
await editorProvider.createNew();
106+
107+
// Check if a temp file is created.
108+
const spiedExecution = sinon.spy(executionService, 'exec');
109+
const spiedWriteFile = sinon.spy(fs, 'writeFile');
110+
disposables.push({ dispose: () => spiedWriteFile.restore() });
111+
disposables.push({ dispose: () => spiedExecution.restore() });
112+
113+
// Format the cell
114+
await commands.executeCommand('notebook.formatCell');
115+
116+
// Verify a temp file was created, having a file starting with thenb file name.
117+
assert.isOk(
118+
spiedWriteFile.getCalls().some((call) => call.args[0].includes('.ipynb')),
119+
'Temp file not created'
120+
);
121+
// Verify we tried to format.
122+
assert.isOk(
123+
spiedExecution.getCalls().some((call) => call.args[0].product === product),
124+
'Not formatted'
125+
);
126+
});
127+
});
128+
});
129+
});

0 commit comments

Comments
 (0)