From b55fc0a9aa9794741d9282c67265d3897e8a0b64 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Sun, 1 Oct 2023 16:43:04 -0400 Subject: [PATCH] Write full Jupyter notebook to stdout --- crates/ruff_cli/src/diagnostics.rs | 4 +- crates/ruff_cli/tests/integration_test.rs | 171 +++++++++++++++++++++- crates/ruff_linter/src/source_kind.rs | 15 ++ crates/ruff_notebook/src/notebook.rs | 6 +- 4 files changed, 190 insertions(+), 6 deletions(-) diff --git a/crates/ruff_cli/src/diagnostics.rs b/crates/ruff_cli/src/diagnostics.rs index 87e35922abc78..fc49b0bfe648b 100644 --- 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 index ae7c840c51f81..94fa4926f1d85 100644 --- 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 index 501d417366613..b065f3e7c8d33 100644 --- 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 index 4d5d0204eeb2a..f1e3e73bffeb9 100644 --- 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())?; }