Skip to content

Commit 4e4b181

Browse files
ArtmannArtmann
committed
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.
1 parent 1851599 commit 4e4b181

File tree

5 files changed

+383
-0
lines changed

5 files changed

+383
-0
lines changed

‎src/notebooks/deepnote/deepnoteNewCellLanguageService.ts‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
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+
// Only process empty code cells
32+
if (cell.kind !== NotebookCellKind.Code) continue;
33+
if (cell.document.getText().trim().length > 0) continue;
34+
35+
// Check if this is an intentional special block (has __deepnotePocket metadata)
36+
const pocketType = cell.metadata?.__deepnotePocket?.type;
37+
if (pocketType) {
38+
// This is an intentional SQL, chart, or input block - keep its language
39+
continue;
40+
}
41+
42+
// If the cell inherited a non-Python language, reset to Python
43+
if (cell.document.languageId !== PYTHON_LANGUAGE) {
44+
languages.setTextDocumentLanguage(cell.document, PYTHON_LANGUAGE).then(noop, noop);
45+
}
46+
}
47+
}
48+
}
49+
}

‎src/notebooks/deepnote/deepnoteNewCellLanguageService.unit.test.ts‎

Lines changed: 323 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,323 @@
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 mockDisposableRegistry: IDisposableRegistry;
13+
let notebookChangeHandler: ((e: any) => void) | undefined;
14+
15+
function createMockNotebook(notebookType: string): NotebookDocument {
16+
return {
17+
uri: Uri.file('/test/notebook.deepnote'),
18+
notebookType
19+
} as NotebookDocument;
20+
}
21+
22+
function createMockCell(options: {
23+
kind?: NotebookCellKind;
24+
languageId?: string;
25+
content?: string;
26+
metadata?: Record<string, unknown>;
27+
}): NotebookCell {
28+
const { kind = NotebookCellKind.Code, languageId = 'python', content = '', metadata = {} } = options;
29+
30+
return {
31+
index: 0,
32+
notebook: createMockNotebook('deepnote'),
33+
kind,
34+
document: {
35+
uri: Uri.file('/test/notebook.deepnote#cell0'),
36+
languageId,
37+
getText: () => content
38+
} as TextDocument,
39+
metadata,
40+
outputs: [],
41+
executionSummary: undefined
42+
} as unknown as NotebookCell;
43+
}
44+
45+
setup(() => {
46+
resetVSCodeMocks();
47+
disposables = [];
48+
mockDisposableRegistry = disposables as unknown as IDisposableRegistry;
49+
50+
// Capture the notebook change handler when workspace.onDidChangeNotebookDocument is called
51+
when(mockedVSCodeNamespaces.workspace.onDidChangeNotebookDocument(anything(), anything())).thenCall(
52+
(handler, thisArg) => {
53+
notebookChangeHandler = (e: any) => handler.call(thisArg, e);
54+
55+
return { dispose: () => undefined };
56+
}
57+
);
58+
59+
// Mock languages.setTextDocumentLanguage to return a resolved promise
60+
when(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(anything(), anything())).thenReturn(
61+
Promise.resolve({} as TextDocument)
62+
);
63+
64+
service = new DeepnoteNewCellLanguageService(mockDisposableRegistry);
65+
});
66+
67+
teardown(() => {
68+
resetVSCodeMocks();
69+
notebookChangeHandler = undefined;
70+
disposables.forEach((d) => d.dispose());
71+
});
72+
73+
suite('activate', () => {
74+
test('registers workspace.onDidChangeNotebookDocument listener', () => {
75+
service.activate();
76+
77+
verify(mockedVSCodeNamespaces.workspace.onDidChangeNotebookDocument(anything(), anything())).once();
78+
});
79+
80+
test('adds disposable to registry', () => {
81+
// Reset mocks to isolate this test
82+
resetVSCodeMocks();
83+
when(mockedVSCodeNamespaces.workspace.onDidChangeNotebookDocument(anything(), anything())).thenCall(
84+
(handler, thisArg) => {
85+
notebookChangeHandler = (e: any) => handler.call(thisArg, e);
86+
87+
return { dispose: () => undefined };
88+
}
89+
);
90+
disposables = [];
91+
mockDisposableRegistry = disposables as unknown as IDisposableRegistry;
92+
service = new DeepnoteNewCellLanguageService(mockDisposableRegistry);
93+
94+
service.activate();
95+
96+
expect(disposables.length).to.be.greaterThan(0);
97+
});
98+
});
99+
100+
suite('onDidChangeNotebookDocument', () => {
101+
setup(() => {
102+
// Reset mock verification state between tests
103+
resetVSCodeMocks();
104+
105+
// Re-setup the mocks after reset
106+
when(mockedVSCodeNamespaces.workspace.onDidChangeNotebookDocument(anything(), anything())).thenCall(
107+
(handler, thisArg) => {
108+
notebookChangeHandler = (e: any) => handler.call(thisArg, e);
109+
110+
return { dispose: () => undefined };
111+
}
112+
);
113+
when(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(anything(), anything())).thenReturn(
114+
Promise.resolve({} as TextDocument)
115+
);
116+
117+
disposables = [];
118+
mockDisposableRegistry = disposables as unknown as IDisposableRegistry;
119+
service = new DeepnoteNewCellLanguageService(mockDisposableRegistry);
120+
service.activate();
121+
expect(notebookChangeHandler).to.not.be.undefined;
122+
});
123+
124+
test('ignores non-deepnote notebooks', async () => {
125+
const jupyterNotebook = createMockNotebook('jupyter-notebook');
126+
const cell = createMockCell({ languageId: 'sql' });
127+
128+
notebookChangeHandler!({
129+
notebook: jupyterNotebook,
130+
contentChanges: [{ addedCells: [cell] }]
131+
});
132+
133+
// Allow async operations to complete
134+
await new Promise((resolve) => setTimeout(resolve, 10));
135+
136+
verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(anything(), anything())).never();
137+
});
138+
139+
test('ignores markdown cells', async () => {
140+
const notebook = createMockNotebook('deepnote');
141+
const cell = createMockCell({ kind: NotebookCellKind.Markup, languageId: 'markdown' });
142+
143+
notebookChangeHandler!({
144+
notebook,
145+
contentChanges: [{ addedCells: [cell] }]
146+
});
147+
148+
await new Promise((resolve) => setTimeout(resolve, 10));
149+
150+
verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(anything(), anything())).never();
151+
});
152+
153+
test('ignores cells with content', async () => {
154+
const notebook = createMockNotebook('deepnote');
155+
const cell = createMockCell({ languageId: 'sql', content: 'SELECT * FROM table' });
156+
157+
notebookChangeHandler!({
158+
notebook,
159+
contentChanges: [{ addedCells: [cell] }]
160+
});
161+
162+
await new Promise((resolve) => setTimeout(resolve, 10));
163+
164+
verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(anything(), anything())).never();
165+
});
166+
167+
test('ignores cells that already have Python language', async () => {
168+
const notebook = createMockNotebook('deepnote');
169+
const cell = createMockCell({ languageId: 'python' });
170+
171+
notebookChangeHandler!({
172+
notebook,
173+
contentChanges: [{ addedCells: [cell] }]
174+
});
175+
176+
await new Promise((resolve) => setTimeout(resolve, 10));
177+
178+
verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(anything(), anything())).never();
179+
});
180+
181+
test('ignores intentional SQL blocks (with __deepnotePocket.type)', async () => {
182+
const notebook = createMockNotebook('deepnote');
183+
const cell = createMockCell({
184+
languageId: 'sql',
185+
metadata: { __deepnotePocket: { type: 'sql' } }
186+
});
187+
188+
notebookChangeHandler!({
189+
notebook,
190+
contentChanges: [{ addedCells: [cell] }]
191+
});
192+
193+
await new Promise((resolve) => setTimeout(resolve, 10));
194+
195+
verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(anything(), anything())).never();
196+
});
197+
198+
test('ignores intentional chart blocks (with __deepnotePocket.type)', async () => {
199+
const notebook = createMockNotebook('deepnote');
200+
const cell = createMockCell({
201+
languageId: 'json',
202+
metadata: { __deepnotePocket: { type: 'chart-vega' } }
203+
});
204+
205+
notebookChangeHandler!({
206+
notebook,
207+
contentChanges: [{ addedCells: [cell] }]
208+
});
209+
210+
await new Promise((resolve) => setTimeout(resolve, 10));
211+
212+
verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(anything(), anything())).never();
213+
});
214+
215+
test('ignores intentional input blocks (with __deepnotePocket.type)', async () => {
216+
const notebook = createMockNotebook('deepnote');
217+
const cell = createMockCell({
218+
languageId: 'plaintext',
219+
metadata: { __deepnotePocket: { type: 'input-text' } }
220+
});
221+
222+
notebookChangeHandler!({
223+
notebook,
224+
contentChanges: [{ addedCells: [cell] }]
225+
});
226+
227+
await new Promise((resolve) => setTimeout(resolve, 10));
228+
229+
verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(anything(), anything())).never();
230+
});
231+
232+
test('changes SQL cell to Python when no __deepnotePocket metadata', async () => {
233+
const notebook = createMockNotebook('deepnote');
234+
const cell = createMockCell({ languageId: 'sql' });
235+
236+
notebookChangeHandler!({
237+
notebook,
238+
contentChanges: [{ addedCells: [cell] }]
239+
});
240+
241+
await new Promise((resolve) => setTimeout(resolve, 10));
242+
243+
verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(cell.document, 'python')).once();
244+
});
245+
246+
test('changes JSON cell to Python when no __deepnotePocket metadata', async () => {
247+
const notebook = createMockNotebook('deepnote');
248+
const cell = createMockCell({ languageId: 'json' });
249+
250+
notebookChangeHandler!({
251+
notebook,
252+
contentChanges: [{ addedCells: [cell] }]
253+
});
254+
255+
await new Promise((resolve) => setTimeout(resolve, 10));
256+
257+
verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(cell.document, 'python')).once();
258+
});
259+
260+
test('handles multiple added cells', async () => {
261+
const notebook = createMockNotebook('deepnote');
262+
const sqlCell = createMockCell({ languageId: 'sql' });
263+
const pythonCell = createMockCell({ languageId: 'python' });
264+
const jsonCell = createMockCell({ languageId: 'json' });
265+
266+
notebookChangeHandler!({
267+
notebook,
268+
contentChanges: [{ addedCells: [sqlCell, pythonCell, jsonCell] }]
269+
});
270+
271+
await new Promise((resolve) => setTimeout(resolve, 10));
272+
273+
// Should change SQL and JSON, but not Python
274+
verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(sqlCell.document, 'python')).once();
275+
verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(pythonCell.document, anything())).never();
276+
verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(jsonCell.document, 'python')).once();
277+
});
278+
279+
test('handles multiple content changes', async () => {
280+
const notebook = createMockNotebook('deepnote');
281+
const cell1 = createMockCell({ languageId: 'sql' });
282+
const cell2 = createMockCell({ languageId: 'javascript' });
283+
284+
notebookChangeHandler!({
285+
notebook,
286+
contentChanges: [{ addedCells: [cell1] }, { addedCells: [cell2] }]
287+
});
288+
289+
await new Promise((resolve) => setTimeout(resolve, 10));
290+
291+
verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(cell1.document, 'python')).once();
292+
verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(cell2.document, 'python')).once();
293+
});
294+
295+
test('ignores content changes with no added cells', async () => {
296+
const notebook = createMockNotebook('deepnote');
297+
298+
notebookChangeHandler!({
299+
notebook,
300+
contentChanges: [{ addedCells: [] }]
301+
});
302+
303+
await new Promise((resolve) => setTimeout(resolve, 10));
304+
305+
verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(anything(), anything())).never();
306+
});
307+
308+
test('ignores cells with whitespace-only content', async () => {
309+
const notebook = createMockNotebook('deepnote');
310+
const cell = createMockCell({ languageId: 'sql', content: ' \n\t ' });
311+
312+
notebookChangeHandler!({
313+
notebook,
314+
contentChanges: [{ addedCells: [cell] }]
315+
});
316+
317+
await new Promise((resolve) => setTimeout(resolve, 10));
318+
319+
// Cell with only whitespace is considered empty, so it should be changed
320+
verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(cell.document, 'python')).once();
321+
});
322+
});
323+
});

‎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);

0 commit comments

Comments
 (0)