Skip to content

Commit 722b206

Browse files
authored
fix: Ensure that Code blocks are Code blocks (#244)
* fix: Ensure that Code blocks are Code blocks. Adding code blocks is handled by the VS Code so we can't modify that behaviour, so instead we check that the block has the correct settings after it's created. * cleanup * clean up the test
1 parent bb8c843 commit 722b206

File tree

5 files changed

+338
-0
lines changed

5 files changed

+338
-0
lines changed
Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,52 @@
1+
import { inject, injectable } from 'inversify';
2+
import { languages, NotebookCellKind, NotebookDocumentChangeEvent, workspace } from 'vscode';
3+
4+
import { DEEPNOTE_NOTEBOOK_TYPE } from '../../kernels/deepnote/types';
5+
import { IExtensionSyncActivationService } from '../../platform/activation/types';
6+
import { PYTHON_LANGUAGE } from '../../platform/common/constants';
7+
import { IDisposableRegistry } from '../../platform/common/types';
8+
import { noop } from '../../platform/common/utils/misc';
9+
10+
/**
11+
* Ensures newly added code cells in Deepnote notebooks default to Python language.
12+
* VS Code copies the language from adjacent cells when inserting, which causes
13+
* new cells after SQL blocks to be SQL. This service corrects that by resetting
14+
* unintentional language inheritance to Python.
15+
*/
16+
@injectable()
17+
export class DeepnoteNewCellLanguageService implements IExtensionSyncActivationService {
18+
constructor(@inject(IDisposableRegistry) private readonly disposables: IDisposableRegistry) {}
19+
20+
public activate(): void {
21+
this.disposables.push(workspace.onDidChangeNotebookDocument(this.onDidChangeNotebookDocument, this));
22+
}
23+
24+
private async onDidChangeNotebookDocument(e: NotebookDocumentChangeEvent): Promise<void> {
25+
if (e.notebook.notebookType !== DEEPNOTE_NOTEBOOK_TYPE) {
26+
return;
27+
}
28+
29+
for (const change of e.contentChanges) {
30+
for (const cell of change.addedCells) {
31+
if (cell.kind !== NotebookCellKind.Code) {
32+
continue;
33+
}
34+
35+
if (cell.document.getText().trim().length > 0) {
36+
continue;
37+
}
38+
39+
const pocketType = cell.metadata?.__deepnotePocket?.type;
40+
41+
if (pocketType) {
42+
continue;
43+
}
44+
45+
// TODO: This will have to be revisited if we add support for other languages in Deepnote.
46+
if (cell.document.languageId !== PYTHON_LANGUAGE) {
47+
languages.setTextDocumentLanguage(cell.document, PYTHON_LANGUAGE).then(noop, noop);
48+
}
49+
}
50+
}
51+
}
52+
}
Lines changed: 275 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,275 @@
1+
import { expect } from 'chai';
2+
import { anything, verify, when } from 'ts-mockito';
3+
import { Disposable, NotebookCell, NotebookCellKind, NotebookDocument, TextDocument, Uri } from 'vscode';
4+
5+
import { IDisposableRegistry } from '../../platform/common/types';
6+
import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../test/vscode-mock';
7+
import { DeepnoteNewCellLanguageService } from './deepnoteNewCellLanguageService';
8+
9+
suite('DeepnoteNewCellLanguageService', () => {
10+
let service: DeepnoteNewCellLanguageService;
11+
let disposables: Disposable[];
12+
let notebookChangeHandler: ((e: any) => void) | undefined;
13+
14+
function createMockNotebook(notebookType: string): NotebookDocument {
15+
return {
16+
uri: Uri.file('/test/notebook.deepnote'),
17+
notebookType
18+
} as NotebookDocument;
19+
}
20+
21+
function createMockCell(options: {
22+
kind?: NotebookCellKind;
23+
languageId?: string;
24+
content?: string;
25+
metadata?: Record<string, unknown>;
26+
}): NotebookCell {
27+
const { kind = NotebookCellKind.Code, languageId = 'python', content = '', metadata = {} } = options;
28+
29+
return {
30+
index: 0,
31+
notebook: createMockNotebook('deepnote'),
32+
kind,
33+
document: {
34+
uri: Uri.file('/test/notebook.deepnote#cell0'),
35+
languageId,
36+
getText: () => content
37+
} as TextDocument,
38+
metadata,
39+
outputs: [],
40+
executionSummary: undefined
41+
} as unknown as NotebookCell;
42+
}
43+
44+
setup(() => {
45+
resetVSCodeMocks();
46+
disposables = [];
47+
48+
when(mockedVSCodeNamespaces.workspace.onDidChangeNotebookDocument(anything(), anything())).thenCall(
49+
(handler, thisArg) => {
50+
notebookChangeHandler = (e: any) => handler.call(thisArg, e);
51+
52+
return { dispose: () => undefined };
53+
}
54+
);
55+
56+
when(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(anything(), anything())).thenReturn(
57+
Promise.resolve({} as TextDocument)
58+
);
59+
60+
service = new DeepnoteNewCellLanguageService(disposables as unknown as IDisposableRegistry);
61+
});
62+
63+
teardown(() => {
64+
notebookChangeHandler = undefined;
65+
disposables.forEach((d) => d.dispose());
66+
});
67+
68+
test('activate registers onDidChangeNotebookDocument listener', () => {
69+
service.activate();
70+
71+
verify(mockedVSCodeNamespaces.workspace.onDidChangeNotebookDocument(anything(), anything())).once();
72+
});
73+
74+
test('activate adds disposable to registry', () => {
75+
service.activate();
76+
77+
expect(disposables.length).to.be.greaterThan(0);
78+
});
79+
80+
test('ignores non-deepnote notebooks', async () => {
81+
service.activate();
82+
const jupyterNotebook = createMockNotebook('jupyter-notebook');
83+
const cell = createMockCell({ languageId: 'sql' });
84+
85+
notebookChangeHandler!({
86+
notebook: jupyterNotebook,
87+
contentChanges: [{ addedCells: [cell] }]
88+
});
89+
await new Promise((resolve) => setTimeout(resolve, 10));
90+
91+
verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(anything(), anything())).never();
92+
});
93+
94+
test('ignores markdown cells', async () => {
95+
service.activate();
96+
const notebook = createMockNotebook('deepnote');
97+
const cell = createMockCell({ kind: NotebookCellKind.Markup, languageId: 'markdown' });
98+
99+
notebookChangeHandler!({
100+
notebook,
101+
contentChanges: [{ addedCells: [cell] }]
102+
});
103+
await new Promise((resolve) => setTimeout(resolve, 10));
104+
105+
verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(anything(), anything())).never();
106+
});
107+
108+
test('ignores cells with content', async () => {
109+
service.activate();
110+
const notebook = createMockNotebook('deepnote');
111+
const cell = createMockCell({ languageId: 'sql', content: 'SELECT * FROM table' });
112+
113+
notebookChangeHandler!({
114+
notebook,
115+
contentChanges: [{ addedCells: [cell] }]
116+
});
117+
await new Promise((resolve) => setTimeout(resolve, 10));
118+
119+
verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(anything(), anything())).never();
120+
});
121+
122+
test('ignores cells that already have Python language', async () => {
123+
service.activate();
124+
const notebook = createMockNotebook('deepnote');
125+
const cell = createMockCell({ languageId: 'python' });
126+
127+
notebookChangeHandler!({
128+
notebook,
129+
contentChanges: [{ addedCells: [cell] }]
130+
});
131+
await new Promise((resolve) => setTimeout(resolve, 10));
132+
133+
verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(anything(), anything())).never();
134+
});
135+
136+
test('ignores intentional SQL blocks (with __deepnotePocket.type)', async () => {
137+
service.activate();
138+
const notebook = createMockNotebook('deepnote');
139+
const cell = createMockCell({
140+
languageId: 'sql',
141+
metadata: { __deepnotePocket: { type: 'sql' } }
142+
});
143+
144+
notebookChangeHandler!({
145+
notebook,
146+
contentChanges: [{ addedCells: [cell] }]
147+
});
148+
await new Promise((resolve) => setTimeout(resolve, 10));
149+
150+
verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(anything(), anything())).never();
151+
});
152+
153+
test('ignores intentional chart blocks (with __deepnotePocket.type)', async () => {
154+
service.activate();
155+
const notebook = createMockNotebook('deepnote');
156+
const cell = createMockCell({
157+
languageId: 'json',
158+
metadata: { __deepnotePocket: { type: 'chart-vega' } }
159+
});
160+
161+
notebookChangeHandler!({
162+
notebook,
163+
contentChanges: [{ addedCells: [cell] }]
164+
});
165+
await new Promise((resolve) => setTimeout(resolve, 10));
166+
167+
verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(anything(), anything())).never();
168+
});
169+
170+
test('ignores intentional input blocks (with __deepnotePocket.type)', async () => {
171+
service.activate();
172+
const notebook = createMockNotebook('deepnote');
173+
const cell = createMockCell({
174+
languageId: 'plaintext',
175+
metadata: { __deepnotePocket: { type: 'input-text' } }
176+
});
177+
178+
notebookChangeHandler!({
179+
notebook,
180+
contentChanges: [{ addedCells: [cell] }]
181+
});
182+
await new Promise((resolve) => setTimeout(resolve, 10));
183+
184+
verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(anything(), anything())).never();
185+
});
186+
187+
test('changes SQL cell to Python when no __deepnotePocket metadata', async () => {
188+
service.activate();
189+
const notebook = createMockNotebook('deepnote');
190+
const cell = createMockCell({ languageId: 'sql' });
191+
192+
notebookChangeHandler!({
193+
notebook,
194+
contentChanges: [{ addedCells: [cell] }]
195+
});
196+
await new Promise((resolve) => setTimeout(resolve, 10));
197+
198+
verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(cell.document, 'python')).once();
199+
});
200+
201+
test('changes JSON cell to Python when no __deepnotePocket metadata', async () => {
202+
service.activate();
203+
const notebook = createMockNotebook('deepnote');
204+
const cell = createMockCell({ languageId: 'json' });
205+
206+
notebookChangeHandler!({
207+
notebook,
208+
contentChanges: [{ addedCells: [cell] }]
209+
});
210+
await new Promise((resolve) => setTimeout(resolve, 10));
211+
212+
verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(cell.document, 'python')).once();
213+
});
214+
215+
test('handles multiple added cells', async () => {
216+
service.activate();
217+
const notebook = createMockNotebook('deepnote');
218+
const sqlCell = createMockCell({ languageId: 'sql' });
219+
const pythonCell = createMockCell({ languageId: 'python' });
220+
const jsonCell = createMockCell({ languageId: 'json' });
221+
222+
notebookChangeHandler!({
223+
notebook,
224+
contentChanges: [{ addedCells: [sqlCell, pythonCell, jsonCell] }]
225+
});
226+
await new Promise((resolve) => setTimeout(resolve, 10));
227+
228+
verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(sqlCell.document, 'python')).once();
229+
verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(pythonCell.document, anything())).never();
230+
verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(jsonCell.document, 'python')).once();
231+
});
232+
233+
test('handles multiple content changes', async () => {
234+
service.activate();
235+
const notebook = createMockNotebook('deepnote');
236+
const cell1 = createMockCell({ languageId: 'sql' });
237+
const cell2 = createMockCell({ languageId: 'javascript' });
238+
239+
notebookChangeHandler!({
240+
notebook,
241+
contentChanges: [{ addedCells: [cell1] }, { addedCells: [cell2] }]
242+
});
243+
await new Promise((resolve) => setTimeout(resolve, 10));
244+
245+
verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(cell1.document, 'python')).once();
246+
verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(cell2.document, 'python')).once();
247+
});
248+
249+
test('ignores content changes with no added cells', async () => {
250+
service.activate();
251+
const notebook = createMockNotebook('deepnote');
252+
253+
notebookChangeHandler!({
254+
notebook,
255+
contentChanges: [{ addedCells: [] }]
256+
});
257+
await new Promise((resolve) => setTimeout(resolve, 10));
258+
259+
verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(anything(), anything())).never();
260+
});
261+
262+
test('changes cells with whitespace-only content to Python', async () => {
263+
service.activate();
264+
const notebook = createMockNotebook('deepnote');
265+
const cell = createMockCell({ languageId: 'sql', content: ' \n\t ' });
266+
267+
notebookChangeHandler!({
268+
notebook,
269+
contentChanges: [{ addedCells: [cell] }]
270+
});
271+
await new Promise((resolve) => setTimeout(resolve, 10));
272+
273+
verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(cell.document, 'python')).once();
274+
});
275+
});

src/notebooks/serviceRegistry.node.ts

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -82,6 +82,7 @@ import { DeepnoteNotebookEnvironmentMapper } from '../kernels/deepnote/environme
8282
import { DeepnoteNotebookCommandListener } from './deepnote/deepnoteNotebookCommandListener';
8383
import { DeepnoteInputBlockCellStatusBarItemProvider } from './deepnote/deepnoteInputBlockCellStatusBarProvider';
8484
import { DeepnoteBigNumberCellStatusBarProvider } from './deepnote/deepnoteBigNumberCellStatusBarProvider';
85+
import { DeepnoteNewCellLanguageService } from './deepnote/deepnoteNewCellLanguageService';
8586
import { SqlIntegrationStartupCodeProvider } from './deepnote/integrations/sqlIntegrationStartupCodeProvider';
8687
import { DeepnoteCellCopyHandler } from './deepnote/deepnoteCellCopyHandler';
8788
import { DeepnoteEnvironmentTreeDataProvider } from '../kernels/deepnote/environments/deepnoteEnvironmentTreeDataProvider.node';
@@ -213,6 +214,10 @@ export function registerTypes(serviceManager: IServiceManager, isDevMode: boolea
213214
IExtensionSyncActivationService,
214215
DeepnoteBigNumberCellStatusBarProvider
215216
);
217+
serviceManager.addSingleton<IExtensionSyncActivationService>(
218+
IExtensionSyncActivationService,
219+
DeepnoteNewCellLanguageService
220+
);
216221

217222
// Deepnote configuration services
218223
serviceManager.addSingleton<DeepnoteEnvironmentStorage>(DeepnoteEnvironmentStorage, DeepnoteEnvironmentStorage);

src/notebooks/serviceRegistry.web.ts

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,7 @@ import {
5151
} from './deepnote/integrations/types';
5252
import { DeepnoteInputBlockCellStatusBarItemProvider } from './deepnote/deepnoteInputBlockCellStatusBarProvider';
5353
import { DeepnoteBigNumberCellStatusBarProvider } from './deepnote/deepnoteBigNumberCellStatusBarProvider';
54+
import { DeepnoteNewCellLanguageService } from './deepnote/deepnoteNewCellLanguageService';
5455
import { SqlCellStatusBarProvider } from './deepnote/sqlCellStatusBarProvider';
5556
import { IntegrationKernelRestartHandler } from './deepnote/integrations/integrationKernelRestartHandler';
5657

@@ -123,6 +124,10 @@ export function registerTypes(serviceManager: IServiceManager, isDevMode: boolea
123124
IExtensionSyncActivationService,
124125
DeepnoteBigNumberCellStatusBarProvider
125126
);
127+
serviceManager.addSingleton<IExtensionSyncActivationService>(
128+
IExtensionSyncActivationService,
129+
DeepnoteNewCellLanguageService
130+
);
126131
serviceManager.addSingleton<IExtensionSyncActivationService>(
127132
IExtensionSyncActivationService,
128133
SqlCellStatusBarProvider

src/test/mocks/vsc/extHostedTypes.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -103,6 +103,7 @@ export namespace vscMockExtHostedTypes {
103103
}
104104
export enum NotebookCellKind {
105105
Markdown = 1,
106+
Markup = 1, // VS Code uses 'Markup' but some older code uses 'Markdown' - both have value 1
106107
Code = 2
107108
}
108109
export enum NotebookCellExecutionState {

0 commit comments

Comments
 (0)