Skip to content

Commit 3fd3b9e

Browse files
joyceerhljoyceerhl
authored
Refactor TensorBoard prompt and import tracking and add tests (microsoft#15073)
1 parent 48abee1 commit 3fd3b9e

9 files changed

Lines changed: 206 additions & 132 deletions

File tree

‎src/client/tensorBoard/serviceRegistry.ts‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,6 @@ import { TensorBoardFileWatcher } from './tensorBoardFileWatcher';
99
import { TensorBoardImportTracker } from './tensorBoardImportTracker';
1010
import { TensorBoardPrompt } from './tensorBoardPrompt';
1111
import { TensorBoardSessionProvider } from './tensorBoardSessionProvider';
12-
import { ITensorBoardImportTracker } from './types';
1312

1413
export function registerTypes(serviceManager: IServiceManager): void {
1514
serviceManager.addSingleton<IExtensionSingleActivationService>(
@@ -19,8 +18,10 @@ export function registerTypes(serviceManager: IServiceManager): void {
1918
serviceManager.addSingleton<TensorBoardFileWatcher>(TensorBoardFileWatcher, TensorBoardFileWatcher);
2019
serviceManager.addBinding(TensorBoardFileWatcher, IExtensionSingleActivationService);
2120
serviceManager.addSingleton<TensorBoardPrompt>(TensorBoardPrompt, TensorBoardPrompt);
22-
serviceManager.addSingleton<ITensorBoardImportTracker>(ITensorBoardImportTracker, TensorBoardImportTracker);
23-
serviceManager.addBinding(ITensorBoardImportTracker, IExtensionSingleActivationService);
21+
serviceManager.addSingleton<IExtensionSingleActivationService>(
22+
IExtensionSingleActivationService,
23+
TensorBoardImportTracker,
24+
);
2425
serviceManager.addSingleton<IExtensionSingleActivationService>(
2526
IExtensionSingleActivationService,
2627
TensorBoardCodeLensProvider,

‎src/client/tensorBoard/tensorBoardImportTracker.ts‎

Lines changed: 5 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -3,36 +3,24 @@
33

44
import { inject, injectable } from 'inversify';
55
import * as path from 'path';
6-
import { Event, EventEmitter, TextEditor } from 'vscode';
6+
import { TextEditor } from 'vscode';
77
import { IExtensionSingleActivationService } from '../activation/types';
88
import { IDocumentManager } from '../common/application/types';
99
import { isTestExecution } from '../common/constants';
1010
import { IDisposableRegistry } from '../common/types';
1111
import { getDocumentLines } from '../telemetry/importTracker';
1212
import { containsTensorBoardImport } from './helpers';
13-
import { ITensorBoardImportTracker } from './types';
13+
import { TensorBoardPrompt } from './tensorBoardPrompt';
1414

1515
const testExecution = isTestExecution();
1616
@injectable()
17-
export class TensorBoardImportTracker implements ITensorBoardImportTracker, IExtensionSingleActivationService {
18-
private pendingChecks = new Map<string, NodeJS.Timer | number>();
19-
20-
private _onDidImportTensorBoard = new EventEmitter<void>();
21-
17+
export class TensorBoardImportTracker implements IExtensionSingleActivationService {
2218
constructor(
2319
@inject(IDocumentManager) private documentManager: IDocumentManager,
2420
@inject(IDisposableRegistry) private disposables: IDisposableRegistry,
21+
@inject(TensorBoardPrompt) private prompt: TensorBoardPrompt,
2522
) {}
2623

27-
// Fires when the active text editor contains a tensorboard import.
28-
public get onDidImportTensorBoard(): Event<void> {
29-
return this._onDidImportTensorBoard.event;
30-
}
31-
32-
public dispose(): void {
33-
this.pendingChecks.clear();
34-
}
35-
3624
public async activate(): Promise<void> {
3725
if (testExecution) {
3826
await this.activateInternal();
@@ -63,7 +51,7 @@ export class TensorBoardImportTracker implements ITensorBoardImportTracker, IExt
6351
) {
6452
const lines = getDocumentLines(document);
6553
if (containsTensorBoardImport(lines)) {
66-
this._onDidImportTensorBoard.fire();
54+
this.prompt.showNativeTensorBoardPrompt().ignoreErrors();
6755
}
6856
}
6957
}

‎src/client/tensorBoard/tensorBoardPrompt.ts‎

Lines changed: 4 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -5,9 +5,8 @@ import { inject, injectable } from 'inversify';
55
import { IApplicationShell, ICommandManager } from '../common/application/types';
66
import { Commands } from '../common/constants';
77
import { NativeTensorBoard } from '../common/experiments/groups';
8-
import { IDisposableRegistry, IExperimentService, IPersistentState, IPersistentStateFactory } from '../common/types';
8+
import { IExperimentService, IPersistentState, IPersistentStateFactory } from '../common/types';
99
import { Common, TensorBoard } from '../common/utils/localize';
10-
import { ITensorBoardImportTracker } from './types';
1110

1211
enum TensorBoardPromptStateKeys {
1312
ShowNativeTensorBoardPrompt = 'showNativeTensorBoardPrompt',
@@ -17,7 +16,7 @@ enum TensorBoardPromptStateKeys {
1716
export class TensorBoardPrompt {
1817
private state: IPersistentState<boolean>;
1918

20-
private enabled: Promise<boolean>;
19+
private enabled: boolean;
2120

2221
private inExperiment: Promise<boolean>;
2322

@@ -28,8 +27,6 @@ export class TensorBoardPrompt {
2827
constructor(
2928
@inject(IApplicationShell) private applicationShell: IApplicationShell,
3029
@inject(ICommandManager) private commandManager: ICommandManager,
31-
@inject(ITensorBoardImportTracker) private importTracker: ITensorBoardImportTracker,
32-
@inject(IDisposableRegistry) private disposableRegistry: IDisposableRegistry,
3330
@inject(IPersistentStateFactory) private persistentStateFactory: IPersistentStateFactory,
3431
@inject(IExperimentService) private experimentService: IExperimentService,
3532
) {
@@ -39,13 +36,12 @@ export class TensorBoardPrompt {
3936
);
4037
this.enabled = this.isPromptEnabled();
4138
this.inExperiment = this.isInExperiment();
42-
this.importTracker.onDidImportTensorBoard(this.showNativeTensorBoardPrompt, this, this.disposableRegistry);
4339
}
4440

4541
public async showNativeTensorBoardPrompt(): Promise<void> {
4642
if (
4743
(await this.inExperiment) &&
48-
(await this.enabled) &&
44+
this.enabled &&
4945
this.enabledInCurrentSession &&
5046
!this.waitingForUserSelection
5147
) {
@@ -73,7 +69,7 @@ export class TensorBoardPrompt {
7369
}
7470
}
7571

76-
private async isPromptEnabled(): Promise<boolean> {
72+
private isPromptEnabled(): boolean {
7773
return this.state.value;
7874
}
7975

‎src/client/tensorBoard/tensorBoardSession.ts‎

Lines changed: 52 additions & 58 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ import { createPromiseFromCancellation } from '../common/cancellation';
1919
import { traceError, traceInfo } from '../common/logger';
2020
import { tensorboardLauncher } from '../common/process/internal/scripts';
2121
import { IProcessServiceFactory, ObservableExecutionResult } from '../common/process/types';
22-
import { IInstaller, InstallerResponse, Product } from '../common/types';
22+
import { IDisposableRegistry, IInstaller, InstallerResponse, Product } from '../common/types';
2323
import { createDeferred, sleep } from '../common/utils/async';
2424
import { TensorBoard } from '../common/utils/localize';
2525
import { IInterpreterService } from '../interpreter/contracts';
@@ -47,6 +47,7 @@ export class TensorBoardSession {
4747
private readonly workspaceService: IWorkspaceService,
4848
private readonly processServiceFactory: IProcessServiceFactory,
4949
private readonly commandManager: ICommandManager,
50+
private readonly disposables: IDisposableRegistry,
5051
) {}
5152

5253
public async initialize(): Promise<void> {
@@ -241,67 +242,60 @@ export class TensorBoardSession {
241242
const webviewPanel = window.createWebviewPanel('tensorBoardSession', 'TensorBoard', ViewColumn.Two, {
242243
enableScripts: true,
243244
});
245+
webviewPanel.webview.html = `<!DOCTYPE html>
246+
<html lang="en">
247+
<head>
248+
<meta charset="UTF-8">
249+
<meta http-equiv="Content-Security-Policy" content="default-src 'unsafe-inline'; frame-src ${this.url} http: https:;">
250+
<meta name="viewport" content="width=device-width, initial-scale=1.0">
251+
<title>TensorBoard</title>
252+
</head>
253+
<body>
254+
<script type="text/javascript">
255+
function resizeFrame() {
256+
var f = window.document.getElementById('vscode-tensorboard-iframe');
257+
if (f) {
258+
f.style.height = window.innerHeight / 0.7 + "px";
259+
f.style.width = window.innerWidth / 0.7 + "px";
260+
}
261+
}
262+
resizeFrame();
263+
window.addEventListener('resize', resizeFrame);
264+
</script>
265+
<iframe
266+
id="vscode-tensorboard-iframe"
267+
class="responsive-iframe"
268+
sandbox="allow-scripts allow-forms allow-same-origin allow-pointer-lock"
269+
src="${this.url}"
270+
frameborder="0"
271+
border="0"
272+
allowfullscreen
273+
></iframe>
274+
<style>
275+
.responsive-iframe {
276+
transform: scale(0.7);
277+
transform-origin: 0 0;
278+
position: absolute;
279+
top: 0;
280+
left: 0;
281+
overflow: hidden;
282+
display: block;
283+
}
284+
</style>
285+
</body>
286+
</html>`;
244287
this.webviewPanel = webviewPanel;
245-
webviewPanel.onDidDispose(() => {
246-
this.webviewPanel = undefined;
247-
// Kill the running TensorBoard session
248-
this.process?.kill();
249-
this.process = undefined;
250-
});
251-
webviewPanel.onDidChangeViewState(() => {
252-
if (webviewPanel.visible) {
253-
this.update();
254-
}
255-
}, null);
288+
this.disposables.push(
289+
webviewPanel.onDidDispose(() => {
290+
this.webviewPanel = undefined;
291+
// Kill the running TensorBoard session
292+
this.process?.kill();
293+
this.process = undefined;
294+
}),
295+
);
256296
return webviewPanel;
257297
}
258298

259-
private update() {
260-
if (this.webviewPanel) {
261-
this.webviewPanel.webview.html = `<!DOCTYPE html>
262-
<html lang="en">
263-
<head>
264-
<meta charset="UTF-8">
265-
<meta http-equiv="Content-Security-Policy" content="default-src 'unsafe-inline'; frame-src ${this.url};">
266-
<meta name="viewport" content="width=device-width, initial-scale=1.0">
267-
<title>TensorBoard</title>
268-
</head>
269-
<body>
270-
<script type="text/javascript">
271-
function resizeFrame() {
272-
var f = window.document.getElementById('vscode-tensorboard-iframe');
273-
if (f) {
274-
f.style.height = window.innerHeight / 0.7 + "px";
275-
f.style.width = window.innerWidth / 0.7 + "px";
276-
}
277-
}
278-
window.addEventListener('resize', resizeFrame);
279-
</script>
280-
<iframe
281-
id="vscode-tensorboard-iframe"
282-
class="responsive-iframe"
283-
sandbox="allow-scripts allow-forms allow-same-origin allow-pointer-lock"
284-
src="${this.url}"
285-
frameborder="0"
286-
border="0"
287-
allowfullscreen
288-
></iframe>
289-
<style>
290-
.responsive-iframe {
291-
transform: scale(0.7);
292-
transform-origin: 0 0;
293-
position: absolute;
294-
top: 0;
295-
left: 0;
296-
overflow: hidden;
297-
display: block;
298-
}
299-
</style>
300-
</body>
301-
</html>`;
302-
}
303-
}
304-
305299
private autopopulateLogDirectoryPath(): string | undefined {
306300
if (this.workspaceService.rootPath) {
307301
return this.workspaceService.rootPath;

‎src/client/tensorBoard/tensorBoardSessionProvider.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,7 @@ export class TensorBoardSessionProvider implements IExtensionSingleActivationSer
5050
this.workspaceService,
5151
this.processServiceFactory,
5252
this.commandManager,
53+
this.disposables,
5354
);
5455
await newSession.initialize();
5556
} catch (e) {

‎src/test/tensorBoard/helpers.ts‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
import * as TypeMoq from 'typemoq';
2+
import { IApplicationShell, ICommandManager } from '../../client/common/application/types';
3+
import { IExperimentService, IPersistentStateFactory } from '../../client/common/types';
4+
import { TensorBoardPrompt } from '../../client/tensorBoard/tensorBoardPrompt';
5+
import { MockState } from '../interpreters/mocks';
6+
7+
export function createTensorBoardPromptWithMocks(): TensorBoardPrompt {
8+
const appShell = TypeMoq.Mock.ofType<IApplicationShell>();
9+
const commandManager = TypeMoq.Mock.ofType<ICommandManager>();
10+
const persistentStateFactory = TypeMoq.Mock.ofType<IPersistentStateFactory>();
11+
const expService = TypeMoq.Mock.ofType<IExperimentService>();
12+
const persistentState = new MockState(true);
13+
persistentStateFactory
14+
.setup((factory) => {
15+
factory.createWorkspacePersistentState(TypeMoq.It.isAny(), TypeMoq.It.isAny());
16+
})
17+
.returns(() => persistentState);
18+
return new TensorBoardPrompt(
19+
appShell.object,
20+
commandManager.object,
21+
persistentStateFactory.object,
22+
expService.object,
23+
);
24+
}

‎src/test/tensorBoard/tensorBoardImportTracker.unit.test.ts‎

Lines changed: 20 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -1,86 +1,79 @@
1+
import { assert } from 'chai';
12
import * as sinon from 'sinon';
23
import { TensorBoardImportTracker } from '../../client/tensorBoard/tensorBoardImportTracker';
4+
import { TensorBoardPrompt } from '../../client/tensorBoard/tensorBoardPrompt';
35
import { MockDocumentManager } from '../startPage/mockDocumentManager';
6+
import { createTensorBoardPromptWithMocks } from './helpers';
47

58
suite('TensorBoard import tracker', () => {
69
let documentManager: MockDocumentManager;
710
let tensorBoardImportTracker: TensorBoardImportTracker;
8-
let onDidImportTensorBoardListener: sinon.SinonExpectation;
11+
let prompt: TensorBoardPrompt;
12+
let showNativeTensorBoardPrompt: sinon.SinonSpy;
913

1014
setup(() => {
1115
documentManager = new MockDocumentManager();
12-
tensorBoardImportTracker = new TensorBoardImportTracker(documentManager, []);
13-
onDidImportTensorBoardListener = sinon.expectation.create('onDidImportTensorBoardListener');
14-
tensorBoardImportTracker.onDidImportTensorBoard(onDidImportTensorBoardListener);
16+
prompt = createTensorBoardPromptWithMocks();
17+
showNativeTensorBoardPrompt = sinon.spy(prompt, 'showNativeTensorBoardPrompt');
18+
tensorBoardImportTracker = new TensorBoardImportTracker(documentManager, [], prompt);
1519
});
1620

1721
test('Simple tensorboard import in Python file', async () => {
1822
const document = documentManager.addDocument('import tensorboard', 'foo.py');
1923
await documentManager.showTextDocument(document);
2024
await tensorBoardImportTracker.activate();
21-
onDidImportTensorBoardListener.once().verify();
25+
assert.ok(showNativeTensorBoardPrompt.calledOnce);
2226
});
2327
test('Simple tensorboard import in Python ipynb', async () => {
2428
const document = documentManager.addDocument('import tensorboard', 'foo.ipynb');
2529
await documentManager.showTextDocument(document);
2630
await tensorBoardImportTracker.activate();
27-
onDidImportTensorBoardListener.once().verify();
31+
assert.ok(showNativeTensorBoardPrompt.calledOnce);
2832
});
2933
test('`from x.y.tensorboard import z` import', async () => {
3034
const document = documentManager.addDocument('from torch.utils.tensorboard import SummaryWriter', 'foo.py');
3135
await documentManager.showTextDocument(document);
3236
await tensorBoardImportTracker.activate();
33-
onDidImportTensorBoardListener.once().verify();
37+
assert.ok(showNativeTensorBoardPrompt.calledOnce);
3438
});
3539
test('`from x.y import tensorboard` import', async () => {
3640
const document = documentManager.addDocument('from torch.utils import tensorboard', 'foo.py');
3741
await documentManager.showTextDocument(document);
3842
await tensorBoardImportTracker.activate();
39-
onDidImportTensorBoardListener.once().verify();
43+
assert.ok(showNativeTensorBoardPrompt.calledOnce);
4044
});
4145
test('`import x, y` import', async () => {
4246
const document = documentManager.addDocument('import tensorboard, tensorflow', 'foo.py');
4347
await documentManager.showTextDocument(document);
4448
await tensorBoardImportTracker.activate();
45-
onDidImportTensorBoardListener.once().verify();
49+
assert.ok(showNativeTensorBoardPrompt.calledOnce);
4650
});
4751
test('`import pkg as _` import', async () => {
4852
const document = documentManager.addDocument('import tensorboard as tb', 'foo.py');
4953
await documentManager.showTextDocument(document);
5054
await tensorBoardImportTracker.activate();
51-
onDidImportTensorBoardListener.once().verify();
55+
assert.ok(showNativeTensorBoardPrompt.calledOnce);
5256
});
53-
test('Fire on changed text editor', async () => {
57+
test('Show prompt on changed text editor', async () => {
5458
await tensorBoardImportTracker.activate();
5559
const document = documentManager.addDocument('import tensorboard as tb', 'foo.py');
5660
await documentManager.showTextDocument(document);
57-
onDidImportTensorBoardListener.once().verify();
61+
assert.ok(showNativeTensorBoardPrompt.calledOnce);
5862
});
59-
test('Do not fire event if no tensorboard import', async () => {
63+
test('Do not show prompt if no tensorboard import', async () => {
6064
const document = documentManager.addDocument('import tensorflow as tf\nfrom torch.utils import foo', 'foo.py');
6165
await documentManager.showTextDocument(document);
6266
await tensorBoardImportTracker.activate();
63-
onDidImportTensorBoardListener.never().verify();
67+
assert.ok(showNativeTensorBoardPrompt.notCalled);
6468
});
65-
test('Do not fire event if language is not Python', async () => {
69+
test('Do not show prompt if language is not Python', async () => {
6670
const document = documentManager.addDocument(
6771
'import tensorflow as tf\nfrom torch.utils import foo',
6872
'foo.cpp',
6973
'cpp',
7074
);
7175
await documentManager.showTextDocument(document);
7276
await tensorBoardImportTracker.activate();
73-
onDidImportTensorBoardListener.never().verify();
74-
});
75-
test('Ignore docstrings', async () => {
76-
const document = documentManager.addDocument(
77-
`"""
78-
import tensorboard
79-
"""`,
80-
'foo.py',
81-
);
82-
await documentManager.showTextDocument(document);
83-
await tensorBoardImportTracker.activate();
84-
onDidImportTensorBoardListener.never().verify();
77+
assert.ok(showNativeTensorBoardPrompt.notCalled);
8578
});
8679
});

0 commit comments

Comments
 (0)