Write full Jupyter notebook to stdout (#7748)

## Summary When writing back notebooks via `stdout`, we need to write back the entire JSON content, not _just_ the fixed source code. Otherwise, writing the output _back_ to the file will yield an invalid notebook. Closes #7747 ## Test Plan `cargo test`
astral-sh · Oct 2, 2023 · ebdfcee · ebdfcee
1 parent c71ff7e
commit ebdfcee
Show file tree

Hide file tree

Showing 4 changed files with 190 additions and 6 deletions.
diff --git a/crates/ruff_cli/src/diagnostics.rs b/crates/ruff_cli/src/diagnostics.rs
@@ -402,7 +402,7 @@ pub(crate) fn lint_stdin(
             match fix_mode {
                 flags::FixMode::Apply => {
                     // Write the contents to stdout, regardless of whether any errors were fixed.
-                    io::stdout().write_all(transformed.source_code().as_bytes())?;
+                    transformed.write(&mut io::stdout().lock())?;
                 }
                 flags::FixMode::Diff => {
                     // But only write a diff if it's non-empty.
@@ -441,7 +441,7 @@ pub(crate) fn lint_stdin(
 
             // Write the contents to stdout anyway.
             if fix_mode.is_apply() {
-                io::stdout().write_all(source_kind.source_code().as_bytes())?;
+                source_kind.write(&mut io::stdout().lock())?;
             }
 
             (result, fixed)

diff --git a/crates/ruff_cli/tests/integration_test.rs b/crates/ruff_cli/tests/integration_test.rs
@@ -75,8 +75,8 @@ fn stdin_filename() {
     "###);
 }
 
-#[test]
 /// Raise `TCH` errors in `.py` files ...
+#[test]
 fn stdin_source_type_py() {
     assert_cmd_snapshot!(Command::new(get_cargo_bin(BIN_NAME))
         .args(STDIN_BASE_OPTIONS)
@@ -136,7 +136,7 @@ fn stdin_json() {
 }
 
 #[test]
-fn stdin_fix() {
+fn stdin_fix_py() {
     let args = ["--fix"];
     assert_cmd_snapshot!(Command::new(get_cargo_bin(BIN_NAME))
         .args(STDIN_BASE_OPTIONS)
@@ -153,6 +153,173 @@ fn stdin_fix() {
     "###);
 }
 
+#[test]
+fn stdin_fix_jupyter() {
+    let args = ["--fix", "--stdin-filename", "Jupyter.ipynb"];
+    assert_cmd_snapshot!(Command::new(get_cargo_bin(BIN_NAME))
+        .args(STDIN_BASE_OPTIONS)
+        .args(args)
+        .pass_stdin(r#"{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "dccc687c-96e2-4604-b957-a8a89b5bec06",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "19e1b029-f516-4662-a9b9-623b93edac1a",
+   "metadata": {},
+   "source": [
+    "Foo"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "cdce7b92-b0fb-4c02-86f6-e233b26fa84f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import sys"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "e40b33d2-7fe4-46c5-bdf0-8802f3052565",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "1\n"
+     ]
+    }
+   ],
+   "source": [
+    "print(1)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a1899bc8-d46f-4ec0-b1d1-e1ca0f04bf60",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.2"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}"#), @r###"
+    success: true
+    exit_code: 0
+    ----- stdout -----
+    {
+     "cells": [
+      {
+       "cell_type": "code",
+       "execution_count": 1,
+       "id": "dccc687c-96e2-4604-b957-a8a89b5bec06",
+       "metadata": {},
+       "outputs": [],
+       "source": []
+      },
+      {
+       "cell_type": "markdown",
+       "id": "19e1b029-f516-4662-a9b9-623b93edac1a",
+       "metadata": {},
+       "source": [
+        "Foo"
+       ]
+      },
+      {
+       "cell_type": "code",
+       "execution_count": 2,
+       "id": "cdce7b92-b0fb-4c02-86f6-e233b26fa84f",
+       "metadata": {},
+       "outputs": [],
+       "source": []
+      },
+      {
+       "cell_type": "code",
+       "execution_count": 3,
+       "id": "e40b33d2-7fe4-46c5-bdf0-8802f3052565",
+       "metadata": {},
+       "outputs": [
+        {
+         "name": "stdout",
+         "output_type": "stream",
+         "text": [
+          "1\n"
+         ]
+        }
+       ],
+       "source": [
+        "print(1)"
+       ]
+      },
+      {
+       "cell_type": "code",
+       "execution_count": null,
+       "id": "a1899bc8-d46f-4ec0-b1d1-e1ca0f04bf60",
+       "metadata": {},
+       "outputs": [],
+       "source": []
+      }
+     ],
+     "metadata": {
+      "kernelspec": {
+       "display_name": "Python 3 (ipykernel)",
+       "language": "python",
+       "name": "python3"
+      },
+      "language_info": {
+       "codemirror_mode": {
+        "name": "ipython",
+        "version": 3
+       },
+       "file_extension": ".py",
+       "mimetype": "text/x-python",
+       "name": "python",
+       "nbconvert_exporter": "python",
+       "pygments_lexer": "ipython3",
+       "version": "3.11.2"
+      }
+     },
+     "nbformat": 4,
+     "nbformat_minor": 5
+    }
+    ----- stderr -----
+    "###);
+}
+
 #[test]
 fn stdin_fix_when_not_fixable_should_still_print_contents() {
     let args = ["--fix"];

diff --git a/crates/ruff_linter/src/source_kind.rs b/crates/ruff_linter/src/source_kind.rs
@@ -1,3 +1,7 @@
+use std::io::Write;
+
+use anyhow::Result;
+
 use ruff_diagnostics::SourceMap;
 use ruff_notebook::Notebook;
 
@@ -22,10 +26,21 @@ impl SourceKind {
         }
     }
 
+    /// Returns the Python source code for this source kind.
     pub fn source_code(&self) -> &str {
         match self {
             SourceKind::Python(source) => source,
             SourceKind::IpyNotebook(notebook) => notebook.source_code(),
         }
     }
+
+    /// Write the transformed source file to the given writer.
+    ///
+    /// For Jupyter notebooks, this will write out the notebook as JSON.
+    pub fn write(&self, writer: &mut dyn Write) -> Result<()> {
+        match self {
+            SourceKind::Python(source) => writer.write_all(source.as_bytes()).map_err(Into::into),
+            SourceKind::IpyNotebook(notebook) => notebook.write(writer).map_err(Into::into),
+        }
+    }
 }
diff --git a/crates/ruff_notebook/src/notebook.rs b/crates/ruff_notebook/src/notebook.rs
@@ -414,11 +414,13 @@ impl Notebook {
     }
 
     /// Write the notebook back to the given [`Write`] implementor.
-    pub fn write(&self, writer: &mut dyn Write) -> anyhow::Result<()> {
+    pub fn write(&self, writer: &mut dyn Write) -> Result<(), NotebookError> {
         // https://github.com/psf/black/blob/69ca0a4c7a365c5f5eea519a90980bab72cab764/src/black/__init__.py#LL1041
         let formatter = serde_json::ser::PrettyFormatter::with_indent(b" ");
         let mut serializer = serde_json::Serializer::with_formatter(writer, formatter);
-        SortAlphabetically(&self.raw).serialize(&mut serializer)?;
+        SortAlphabetically(&self.raw)
+            .serialize(&mut serializer)
+            .map_err(NotebookError::Json)?;
         if self.trailing_newline {
             writeln!(serializer.into_inner())?;
         }