VS Code Notebooks with Langium

With Jupyter Notebooks, data scientists and developers using Python have had a de facto standard way of sharing documentation and executable code samples for quite awhile now. In 2021, VS Code rolled out their new notebook experience for everyone. Consequently, developers started using this powerful new API to build more and more notebook extensions.

In this blog post, I want to demonstrate how developers using Langium can leverage VS Code and the Language Server Protocol (LSP) to provide notebook support for their languages. This will be exemplified using the Langium implementation of the Lox language. To add, this implementation of Lox implements a statically typed dialect of the Lox language known from the book Crafting Interpreters.

Language kernel

The first task we want to accomplish is to execute code cells which contain our language code. To do this, we need to register a notebooks contribution in the package.json of our VS Code extension:

{
    "contributes": {
        ...
        "notebooks": [{
            "type": "lox-notebook",
            "displayName": "Lox Notebook",
            "selector": [
                {
                    "filenamePattern": "*.loxnb"
                }
            ]
        }]
    }
}

Now we need to register a new notebook controller for this lox-notebook label.

import * as vscode from 'vscode';

export class LoxNotebookKernel {
    readonly id = 'lox-kernel';
    public readonly label = 'Lox Kernel';
    readonly supportedLanguages = ['lox'];

    private readonly _controller: vscode.NotebookController;

    constructor() {

        this._controller = vscode.notebooks.createNotebookController(this.id,
            'lox-notebook',
            this.label);

        this._controller.supportedLanguages = this.supportedLanguages;
        this._controller.supportsExecutionOrder = true;
    }

    dispose(): void {
        this._controller.dispose();
    }
}

For now, this kernel doesn’t do anything. To rectify this, we need to implement an interpreter function for the Lox language, and call it for each notebook cell (ours will be called _doExecution). We can then import runInterpreter from interpreter/runner.ts. Afterwards, we can bind our interpreter function _doExecution to the executeHandler of our notebook controller.

import { runInterpreter } from '../interpreter/runner';

export class LoxNotebookKernel {
    constructor() {
        ...
        this._controller.executeHandler = this._executeAll.bind(this);
    }

    private async _executeAll(cells: vscode.NotebookCell[]): Promise<void> {
        for (let cell of cells) {
            await this._doExecution(cell);
        }
    }

    private async _doExecution(cell: vscode.NotebookCell): Promise<void> {
        const execution = this._controller.createNotebookCellExecution(cell);

        execution.executionOrder = ++this._executionOrder;
        execution.start(Date.now());

        const text = cell.document.getText();
        await execution.clearOutput();
        const log = async (value: unknown) => {
            const stringValue = `${value}`;
            await execution.appendOutput(
                new vscode.NotebookCellOutput([vscode.NotebookCellOutputItem.text(stringValue)])
            );
        }

        try {
            await runInterpreter(text, { log });
            execution.end(true, Date.now());
        } catch (err) {
            const errString = err instanceof Error ? err.message : String(err);
            await execution.appendOutput(
                new vscode.NotebookCellOutput([vscode.NotebookCellOutputItem.text(errString)])
            );
            execution.end(false, Date.now());
        }
    }
}

VS Code also requires developers to register a serializer for their notebook data. We will use a simple sample serializer for this task. For reference, you can see the serializer for Lox here.

We can then register the kernel and serializer, and push both of them to the subscriptions property of the extension context. This will ensure they’re disposed of when our extension unloads.

export function activate(context: vscode.ExtensionContext): void {
    context.subscriptions.push(
        vscode.workspace.registerNotebookSerializer(
            'lox-notebook', new LoxNotebookSerializer(), { transientOutputs: true }
        ),
        new LoxNotebookKernel()
    );
}

Editing support

At this point, we have notebook support for our language. However, we’re still missing editing support. When you setup a project with Langium, all prerequisites for editor support in a normal editor tab are already set up for you. While syntax highlighting through textMate grammars will still work as expected in notebooks, advanced editor support – such as code completion and diagnostics – will be absent.

Luckily, VS Code has made it very easy to add advanced editing support for notebooks through the LSP. Your language client setup can generally achieve this by listening for document changes under the file URI schema:

const clientOptions: LanguageClientOptions = {
    documentSelector: [{
        scheme: "file",
        language: "lox"
    }]
};

Lastly, all we need to do is add the new vscode-notebook-cell schema, and vscode will start sending the cell contents to the language server. In Langium, each cell will be handled as a separate document, which enables all your LSP features to work in notebook cells.

const clientOptions: LanguageClientOptions = {
    documentSelector: [
        {
            scheme: "file",
            language: "lox"
        }, {
            scheme: 'vscode-notebook-cell',
            language: 'lox'
        }
    ]
};

Conclusion

That’s all you need to create a powerful and integrated notebook experience for any language built with Langium. Of course, the code can also be used as a blueprint to provide notebook support for any kind of language server with interpreter support.

If you’re looking for references or samples related to notebook integration, you can check out the langium-lox repository.