diff --git a/src/notebooks/deepnote/deepnoteNewCellLanguageService.ts b/src/notebooks/deepnote/deepnoteNewCellLanguageService.ts new file mode 100644 index 0000000000..b3844e494d --- /dev/null +++ b/src/notebooks/deepnote/deepnoteNewCellLanguageService.ts @@ -0,0 +1,52 @@ +import { inject, injectable } from 'inversify'; +import { languages, NotebookCellKind, NotebookDocumentChangeEvent, workspace } from 'vscode'; + +import { DEEPNOTE_NOTEBOOK_TYPE } from '../../kernels/deepnote/types'; +import { IExtensionSyncActivationService } from '../../platform/activation/types'; +import { PYTHON_LANGUAGE } from '../../platform/common/constants'; +import { IDisposableRegistry } from '../../platform/common/types'; +import { noop } from '../../platform/common/utils/misc'; + +/** + * Ensures newly added code cells in Deepnote notebooks default to Python language. + * VS Code copies the language from adjacent cells when inserting, which causes + * new cells after SQL blocks to be SQL. This service corrects that by resetting + * unintentional language inheritance to Python. + */ +@injectable() +export class DeepnoteNewCellLanguageService implements IExtensionSyncActivationService { + constructor(@inject(IDisposableRegistry) private readonly disposables: IDisposableRegistry) {} + + public activate(): void { + this.disposables.push(workspace.onDidChangeNotebookDocument(this.onDidChangeNotebookDocument, this)); + } + + private async onDidChangeNotebookDocument(e: NotebookDocumentChangeEvent): Promise { + if (e.notebook.notebookType !== DEEPNOTE_NOTEBOOK_TYPE) { + return; + } + + for (const change of e.contentChanges) { + for (const cell of change.addedCells) { + if (cell.kind !== NotebookCellKind.Code) { + continue; + } + + if (cell.document.getText().trim().length > 0) { + continue; + } + + const pocketType = cell.metadata?.__deepnotePocket?.type; + + if (pocketType) { + continue; + } + + // TODO: This will have to be revisited if we add support for other languages in Deepnote. + if (cell.document.languageId !== PYTHON_LANGUAGE) { + languages.setTextDocumentLanguage(cell.document, PYTHON_LANGUAGE).then(noop, noop); + } + } + } + } +} diff --git a/src/notebooks/deepnote/deepnoteNewCellLanguageService.unit.test.ts b/src/notebooks/deepnote/deepnoteNewCellLanguageService.unit.test.ts new file mode 100644 index 0000000000..6293c9efa2 --- /dev/null +++ b/src/notebooks/deepnote/deepnoteNewCellLanguageService.unit.test.ts @@ -0,0 +1,275 @@ +import { expect } from 'chai'; +import { anything, verify, when } from 'ts-mockito'; +import { Disposable, NotebookCell, NotebookCellKind, NotebookDocument, TextDocument, Uri } from 'vscode'; + +import { IDisposableRegistry } from '../../platform/common/types'; +import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../test/vscode-mock'; +import { DeepnoteNewCellLanguageService } from './deepnoteNewCellLanguageService'; + +suite('DeepnoteNewCellLanguageService', () => { + let service: DeepnoteNewCellLanguageService; + let disposables: Disposable[]; + let notebookChangeHandler: ((e: any) => void) | undefined; + + function createMockNotebook(notebookType: string): NotebookDocument { + return { + uri: Uri.file('/test/notebook.deepnote'), + notebookType + } as NotebookDocument; + } + + function createMockCell(options: { + kind?: NotebookCellKind; + languageId?: string; + content?: string; + metadata?: Record; + }): NotebookCell { + const { kind = NotebookCellKind.Code, languageId = 'python', content = '', metadata = {} } = options; + + return { + index: 0, + notebook: createMockNotebook('deepnote'), + kind, + document: { + uri: Uri.file('/test/notebook.deepnote#cell0'), + languageId, + getText: () => content + } as TextDocument, + metadata, + outputs: [], + executionSummary: undefined + } as unknown as NotebookCell; + } + + setup(() => { + resetVSCodeMocks(); + disposables = []; + + when(mockedVSCodeNamespaces.workspace.onDidChangeNotebookDocument(anything(), anything())).thenCall( + (handler, thisArg) => { + notebookChangeHandler = (e: any) => handler.call(thisArg, e); + + return { dispose: () => undefined }; + } + ); + + when(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(anything(), anything())).thenReturn( + Promise.resolve({} as TextDocument) + ); + + service = new DeepnoteNewCellLanguageService(disposables as unknown as IDisposableRegistry); + }); + + teardown(() => { + notebookChangeHandler = undefined; + disposables.forEach((d) => d.dispose()); + }); + + test('activate registers onDidChangeNotebookDocument listener', () => { + service.activate(); + + verify(mockedVSCodeNamespaces.workspace.onDidChangeNotebookDocument(anything(), anything())).once(); + }); + + test('activate adds disposable to registry', () => { + service.activate(); + + expect(disposables.length).to.be.greaterThan(0); + }); + + test('ignores non-deepnote notebooks', async () => { + service.activate(); + const jupyterNotebook = createMockNotebook('jupyter-notebook'); + const cell = createMockCell({ languageId: 'sql' }); + + notebookChangeHandler!({ + notebook: jupyterNotebook, + contentChanges: [{ addedCells: [cell] }] + }); + await new Promise((resolve) => setTimeout(resolve, 10)); + + verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(anything(), anything())).never(); + }); + + test('ignores markdown cells', async () => { + service.activate(); + const notebook = createMockNotebook('deepnote'); + const cell = createMockCell({ kind: NotebookCellKind.Markup, languageId: 'markdown' }); + + notebookChangeHandler!({ + notebook, + contentChanges: [{ addedCells: [cell] }] + }); + await new Promise((resolve) => setTimeout(resolve, 10)); + + verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(anything(), anything())).never(); + }); + + test('ignores cells with content', async () => { + service.activate(); + const notebook = createMockNotebook('deepnote'); + const cell = createMockCell({ languageId: 'sql', content: 'SELECT * FROM table' }); + + notebookChangeHandler!({ + notebook, + contentChanges: [{ addedCells: [cell] }] + }); + await new Promise((resolve) => setTimeout(resolve, 10)); + + verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(anything(), anything())).never(); + }); + + test('ignores cells that already have Python language', async () => { + service.activate(); + const notebook = createMockNotebook('deepnote'); + const cell = createMockCell({ languageId: 'python' }); + + notebookChangeHandler!({ + notebook, + contentChanges: [{ addedCells: [cell] }] + }); + await new Promise((resolve) => setTimeout(resolve, 10)); + + verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(anything(), anything())).never(); + }); + + test('ignores intentional SQL blocks (with __deepnotePocket.type)', async () => { + service.activate(); + const notebook = createMockNotebook('deepnote'); + const cell = createMockCell({ + languageId: 'sql', + metadata: { __deepnotePocket: { type: 'sql' } } + }); + + notebookChangeHandler!({ + notebook, + contentChanges: [{ addedCells: [cell] }] + }); + await new Promise((resolve) => setTimeout(resolve, 10)); + + verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(anything(), anything())).never(); + }); + + test('ignores intentional chart blocks (with __deepnotePocket.type)', async () => { + service.activate(); + const notebook = createMockNotebook('deepnote'); + const cell = createMockCell({ + languageId: 'json', + metadata: { __deepnotePocket: { type: 'chart-vega' } } + }); + + notebookChangeHandler!({ + notebook, + contentChanges: [{ addedCells: [cell] }] + }); + await new Promise((resolve) => setTimeout(resolve, 10)); + + verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(anything(), anything())).never(); + }); + + test('ignores intentional input blocks (with __deepnotePocket.type)', async () => { + service.activate(); + const notebook = createMockNotebook('deepnote'); + const cell = createMockCell({ + languageId: 'plaintext', + metadata: { __deepnotePocket: { type: 'input-text' } } + }); + + notebookChangeHandler!({ + notebook, + contentChanges: [{ addedCells: [cell] }] + }); + await new Promise((resolve) => setTimeout(resolve, 10)); + + verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(anything(), anything())).never(); + }); + + test('changes SQL cell to Python when no __deepnotePocket metadata', async () => { + service.activate(); + const notebook = createMockNotebook('deepnote'); + const cell = createMockCell({ languageId: 'sql' }); + + notebookChangeHandler!({ + notebook, + contentChanges: [{ addedCells: [cell] }] + }); + await new Promise((resolve) => setTimeout(resolve, 10)); + + verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(cell.document, 'python')).once(); + }); + + test('changes JSON cell to Python when no __deepnotePocket metadata', async () => { + service.activate(); + const notebook = createMockNotebook('deepnote'); + const cell = createMockCell({ languageId: 'json' }); + + notebookChangeHandler!({ + notebook, + contentChanges: [{ addedCells: [cell] }] + }); + await new Promise((resolve) => setTimeout(resolve, 10)); + + verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(cell.document, 'python')).once(); + }); + + test('handles multiple added cells', async () => { + service.activate(); + const notebook = createMockNotebook('deepnote'); + const sqlCell = createMockCell({ languageId: 'sql' }); + const pythonCell = createMockCell({ languageId: 'python' }); + const jsonCell = createMockCell({ languageId: 'json' }); + + notebookChangeHandler!({ + notebook, + contentChanges: [{ addedCells: [sqlCell, pythonCell, jsonCell] }] + }); + await new Promise((resolve) => setTimeout(resolve, 10)); + + verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(sqlCell.document, 'python')).once(); + verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(pythonCell.document, anything())).never(); + verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(jsonCell.document, 'python')).once(); + }); + + test('handles multiple content changes', async () => { + service.activate(); + const notebook = createMockNotebook('deepnote'); + const cell1 = createMockCell({ languageId: 'sql' }); + const cell2 = createMockCell({ languageId: 'javascript' }); + + notebookChangeHandler!({ + notebook, + contentChanges: [{ addedCells: [cell1] }, { addedCells: [cell2] }] + }); + await new Promise((resolve) => setTimeout(resolve, 10)); + + verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(cell1.document, 'python')).once(); + verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(cell2.document, 'python')).once(); + }); + + test('ignores content changes with no added cells', async () => { + service.activate(); + const notebook = createMockNotebook('deepnote'); + + notebookChangeHandler!({ + notebook, + contentChanges: [{ addedCells: [] }] + }); + await new Promise((resolve) => setTimeout(resolve, 10)); + + verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(anything(), anything())).never(); + }); + + test('changes cells with whitespace-only content to Python', async () => { + service.activate(); + const notebook = createMockNotebook('deepnote'); + const cell = createMockCell({ languageId: 'sql', content: ' \n\t ' }); + + notebookChangeHandler!({ + notebook, + contentChanges: [{ addedCells: [cell] }] + }); + await new Promise((resolve) => setTimeout(resolve, 10)); + + verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(cell.document, 'python')).once(); + }); +}); diff --git a/src/notebooks/serviceRegistry.node.ts b/src/notebooks/serviceRegistry.node.ts index a23f0a1c51..8d44992a43 100644 --- a/src/notebooks/serviceRegistry.node.ts +++ b/src/notebooks/serviceRegistry.node.ts @@ -82,6 +82,7 @@ import { DeepnoteNotebookEnvironmentMapper } from '../kernels/deepnote/environme import { DeepnoteNotebookCommandListener } from './deepnote/deepnoteNotebookCommandListener'; import { DeepnoteInputBlockCellStatusBarItemProvider } from './deepnote/deepnoteInputBlockCellStatusBarProvider'; import { DeepnoteBigNumberCellStatusBarProvider } from './deepnote/deepnoteBigNumberCellStatusBarProvider'; +import { DeepnoteNewCellLanguageService } from './deepnote/deepnoteNewCellLanguageService'; import { SqlIntegrationStartupCodeProvider } from './deepnote/integrations/sqlIntegrationStartupCodeProvider'; import { DeepnoteCellCopyHandler } from './deepnote/deepnoteCellCopyHandler'; import { DeepnoteEnvironmentTreeDataProvider } from '../kernels/deepnote/environments/deepnoteEnvironmentTreeDataProvider.node'; @@ -213,6 +214,10 @@ export function registerTypes(serviceManager: IServiceManager, isDevMode: boolea IExtensionSyncActivationService, DeepnoteBigNumberCellStatusBarProvider ); + serviceManager.addSingleton( + IExtensionSyncActivationService, + DeepnoteNewCellLanguageService + ); // Deepnote configuration services serviceManager.addSingleton(DeepnoteEnvironmentStorage, DeepnoteEnvironmentStorage); diff --git a/src/notebooks/serviceRegistry.web.ts b/src/notebooks/serviceRegistry.web.ts index 49f5077030..708e657875 100644 --- a/src/notebooks/serviceRegistry.web.ts +++ b/src/notebooks/serviceRegistry.web.ts @@ -51,6 +51,7 @@ import { } from './deepnote/integrations/types'; import { DeepnoteInputBlockCellStatusBarItemProvider } from './deepnote/deepnoteInputBlockCellStatusBarProvider'; import { DeepnoteBigNumberCellStatusBarProvider } from './deepnote/deepnoteBigNumberCellStatusBarProvider'; +import { DeepnoteNewCellLanguageService } from './deepnote/deepnoteNewCellLanguageService'; import { SqlCellStatusBarProvider } from './deepnote/sqlCellStatusBarProvider'; import { IntegrationKernelRestartHandler } from './deepnote/integrations/integrationKernelRestartHandler'; @@ -123,6 +124,10 @@ export function registerTypes(serviceManager: IServiceManager, isDevMode: boolea IExtensionSyncActivationService, DeepnoteBigNumberCellStatusBarProvider ); + serviceManager.addSingleton( + IExtensionSyncActivationService, + DeepnoteNewCellLanguageService + ); serviceManager.addSingleton( IExtensionSyncActivationService, SqlCellStatusBarProvider diff --git a/src/test/mocks/vsc/extHostedTypes.ts b/src/test/mocks/vsc/extHostedTypes.ts index 962d75bd0e..4a58f1c077 100644 --- a/src/test/mocks/vsc/extHostedTypes.ts +++ b/src/test/mocks/vsc/extHostedTypes.ts @@ -103,6 +103,7 @@ export namespace vscMockExtHostedTypes { } export enum NotebookCellKind { Markdown = 1, + Markup = 1, // VS Code uses 'Markup' but some older code uses 'Markdown' - both have value 1 Code = 2 } export enum NotebookCellExecutionState {