Skip to content

Commit

Permalink
Add support for Jupyter Notebook (#264)
Browse files Browse the repository at this point in the history
## Summary

This PR adds support for Jupyter Notebook. It requires client support
for LSP 3.17 which contains the [Notebook
support](https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#notebookDocument_synchronization).

### Implementation

#### Context

* `Document`: LSP type representing a text file (Python file for Ruff).
* `TextDocument`: `pygls` representation of the LSP `Document`. This is
an abstraction created from a `Document` which provides some useful
methods like getting the file path, source code, etc.
* New in 3.17: `NotebookDocument` type was added to represent a Notebook
which consists of a list of cells (`NotebookCell`). Note that these are
all LSP types coming from `lsprotocol`.
* In `pygls`, a Notebook cell is represented as a text document
(`TextDocument`).

There are methods provided by `pygls` to get the object:
* `get_text_document` - Returns a `TextDocument` which either represents
a Python file or a Notebook cell
* `get_notebook_document` - Returns a `NotebookDocument` either using
the Notebook URI or a cell URI. For cell URI, it returns the
`NotebookDocument` containing the cell.

#### Document

A new `Document` type was created to facilitate the implementation. This
represents either a Python file, a Notebook or a Notebook cell. There
are various constructor methods which should be used to create this
type:
* For a URI representing a Python file, use either `from_uri` or
`from_text_document`.
* For a URI representing a Notebook file, use either `from_uri` or
`from_notebook_document`.
* For a URI representing a Notebook cell, use either
`from_cell_or_text_uri` or `from_notebook_cell`.

#### Notebook JSON

Ruff expects the source content of a Notebook file to be in JSON format
following the [Notebook format specification] but the protocol uses it's
own abstraction and doesn't store the JSON format. This means that we
need to create a JSON string representing the Notebook from the
available information. This doesn't need all the information as Ruff
only uses the cell source and version information. So, we create a
minimal JSON string representing the Notebook document and pass it to
Ruff.

<details><summary>An example JSON string representing a Notebook
Document:</summary>
<p>

```json
{
  "metadata": {},
  "nbformat": 4,
  "nbformat_minor": 5,
  "cells": [
    {
      "cell_type": "code",
      "metadata": null,
      "outputs": [],
      "source": "import random\nimport math"
    },
    {
      "cell_type": "code",
      "metadata": null,
      "outputs": [],
      "source": "try:\n    print()\nexcept ValueError:\n    pass"
    },
    {
      "cell_type": "code",
      "metadata": null,
      "outputs": [],
      "source": "import random\nimport pprint\n\nrandom.randint(10, 20)"
    },
    {
      "cell_type": "code",
      "metadata": null,
      "outputs": [],
      "source": "foo = 1\nif foo is 1:\n    msg = f\"Invalid foo: {foo}\"\n    raise ValueError(msg)"
    }
  ]
}
```

</p>
</details> 

**We need to pass in every cell including the markdown cell to get an
accurate information like the cell number.**

For the cell document kind, the source value is a JSON string containing
just a single code cell. This is required as code actions and formatting
work at both cell and notebook level.

### Configuration

For VSCode users, the `notebook.*` configuration is used to run the
formatter or code actions on save:

```jsonc
{
  // Enable formatting the entire Notebook on save
  "notebook.formatOnSave.enabled": true,
  // Run the enabled code actions on the entire Notebook on save
  "notebook.codeActionsOnSave": {
    "source.fixAll": true,
    "source.organizeImports.ruff": true
  },
}
```

The way the above settings work in VSCode is that the editor runs the
actions in parallel for every cell. This has the illusion that it was
run on the entire Notebook. The commands defined by us (`Ruff: Organize
imports` and `Ruff: Fix all auto-fixable problems`) are run on the
entire Notebook at once. This is important because in the latter case
the `ruff` command is invoked `n` number of times where `n` is the
number of cells while for the former it's run only once.

### Commands

#### Builtin

* `Ruff: Organize Imports`: Works at Notebook level
* `Ruff: Fix all auto-fixable problems`: Works at Notebook level

#### VSCode specifc

* `Format Cell`: Formats the current cell
* `Notebook: Format Notebook`: Formats the entire Notebook by running
the formatter for every cell
* `Organize Imports`: Runs the `source.organizeImports` code action on
every cell in parallel
* `Fix All`: Runs the `source.fixAll` code action on every cell in
parallel

## Feature checklist

- [x] Code actions
  - [x] Organize imports
  - [x] Fix all
  - [x] Each fixable diagnostics
  - [x] Disable rule comment
- [x] Code action resolve
- [x] Commands
  - [x] `ruff.applyAutofix`
  - [x] `ruff.applyOrganizeImports`
  - [x] `ruff.applyFormat`
- [x] Diagnostics
  - [x] On open
  - [x] On close
  - [x] On save
  - [x] On change
- [x] Formatting
- [x] Hover

## Test Plan

Manually testing for all the features mentioned above.

### How to run this locally?

1. Clone https://github.com/astral-sh/ruff-lsp and
https://github.com/astral-sh/ruff-vscode in the same directory
2. Checkout this branch `git checkout dhruv/notebook` in the `ruff-lsp`
repository
3. Install the requirements for both repositories
4. For `ruff-vscode`, uninstall `ruff-lsp` (`pip uninstall --yes
ruff-lsp`) as we'd want to use the local version. To install the local
`ruff-lsp` version in `ruff-vscode`, follow [Modifying the
LSP](https://github.com/astral-sh/ruff-vscode#modifying-the-lsp).
5. Open VSCode from `ruff-vscode` directory -> "Run and Debug" section
from the sidebar -> "Debug Extension and Python" config.

This will then open a VSCode development session which can be used to
test out the notebook features.

**Test notebooks:**
* Formatting:
https://gist.github.com/dhruvmanila/7803e5a3b98c414505384db415a635a0
* Diagnostics, Code actions, Commands:
https://gist.github.com/dhruvmanila/54c65870f167a56558d4701f57f53042

**Requires: astral-sh/ruff#7664 which was
released in `v0.1.0`**

fixes: #267 
closes: astral-sh/ruff-vscode#256
closes: astral-sh/ruff-vscode#314
closes: astral-sh/ruff-vscode#51
  • Loading branch information
dhruvmanila authored Nov 8, 2023
1 parent 3dc7888 commit 31f3836
Show file tree
Hide file tree
Showing 4 changed files with 605 additions and 136 deletions.
Loading

0 comments on commit 31f3836

Please sign in to comment.