Add roundtrip support for Jupyter notebook (#5028)

## Summary

Add roundtrip support for Jupyter notebook.

1. Read the notebook
2. Extract out the source code content
3. Use it to update the notebook itself (should be exactly the same [^1])
4. Serialize into JSON and print it to stdout

## Test Plan

`cargo run --all-features --bin ruff_dev --package ruff_dev --
round-trip <path/to/notebook.ipynb>`

<details><summary>Example output:</summary>
<p>

```
{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f3c286e9-fa52-4440-816f-4449232f199a",
   "metadata": {},
   "source": [
    "# Ruff Test"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a2b7bc6c-778a-4b07-86ae-dde5a2d9511e",
   "metadata": {},
   "source": [
    "Markdown block before the first import"
   ]
  },
  {
   "cell_type": "code",
   "id": "5e3ef98e-224c-450a-80e6-be442ad50907",
   "metadata": {
    "tags": []
   },
   "source": "",
   "execution_count": 1,
   "outputs": []
  },
  {
   "cell_type": "code",
   "id": "6bced3f8-e0a4-450c-ae7c-f60ad5671ee9",
   "metadata": {},
   "source": "import contextlib\n\nwith contextlib.suppress(ValueError):\n    print()\n",
   "outputs": []
  },
  {
   "cell_type": "code",
   "id": "d7102cfd-5bb5-4f5b-a3b8-07a7b8cca34c",
   "metadata": {},
   "source": "import random\n\nrandom.randint(10, 20)",
   "outputs": []
  },
  {
   "cell_type": "code",
   "id": "88471d1c-7429-4967-898f-b0088fcb4c53",
   "metadata": {},
   "source": "foo = 1\nif foo < 2:\n    msg = f\"Invalid foo: {foo}\"\n    raise ValueError(msg)",
   "outputs": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (ruff-playground)",
   "name": "ruff-playground",
   "language": "python"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "pygments_lexer": "ipython3",
   "nbconvert_exporter": "python",
   "version": "3.11.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
```

</p>
</details> 

[^1]: The type in JSON might be different (https://github.com/astral-sh/ruff/pull/4665#discussion_r1212663495)

Part of #1218
This commit is contained in:
Dhruv Manilawala 2023-06-12 23:27:45 +05:30 committed by GitHub
parent a77d2df934
commit cb4f086cbf
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 36 additions and 9 deletions

View file

@ -1,4 +1,4 @@
//! Run round-trip source code generation on a given Python file.
//! Run round-trip source code generation on a given Python or Jupyter notebook file.
#![allow(clippy::print_stdout, clippy::print_stderr)]
use std::fs;
@ -6,17 +6,23 @@ use std::path::PathBuf;
use anyhow::Result;
use ruff::jupyter;
use ruff::round_trip;
#[derive(clap::Args)]
pub(crate) struct Args {
/// Python file to round-trip.
/// Python or Jupyter notebook file to round-trip.
#[arg(required = true)]
file: PathBuf,
}
pub(crate) fn main(args: &Args) -> Result<()> {
let contents = fs::read_to_string(&args.file)?;
println!("{}", round_trip(&contents, &args.file.to_string_lossy())?);
let path = args.file.as_path();
if jupyter::is_jupyter_notebook(path) {
println!("{}", jupyter::round_trip(path)?);
} else {
let contents = fs::read_to_string(&args.file)?;
println!("{}", round_trip(&contents, &args.file.to_string_lossy())?);
}
Ok(())
}