Use transformed source code for diagnostic locations (#9408)

## Summary After we apply fixes, the source code might be transformed. And yet, we're using the _unmodified_ source code to compute locations in some cases (e.g., for displaying parse errors, or Jupyter Notebook cells). This can lead to subtle errors in reporting, or even panics. This PR modifies the linter to use the _transformed_ source code for such computations. Closes #9407.
astral-sh · Jan 6, 2024 · cee0976 · cee0976
1 parent cde4a7d
commit cee0976
Show file tree

Hide file tree

Showing 3 changed files with 44 additions and 23 deletions.
diff --git a/crates/ruff_cli/src/diagnostics.rs b/crates/ruff_cli/src/diagnostics.rs
@@ -1,5 +1,6 @@
 #![cfg_attr(target_family = "wasm", allow(dead_code))]
 
+use std::borrow::Cow;
 use std::fs::File;
 use std::io;
 use std::ops::{Add, AddAssign};
@@ -273,6 +274,7 @@ pub(crate) fn lint_path(
             data: (messages, imports),
             error: parse_error,
         },
+        transformed,
         fixed,
     ) = if matches!(fix_mode, flags::FixMode::Apply | flags::FixMode::Diff) {
         if let Ok(FixerResult {
@@ -301,7 +303,12 @@ pub(crate) fn lint_path(
                     flags::FixMode::Generate => {}
                 }
             }
-            (result, fixed)
+            let transformed = if let Cow::Owned(transformed) = transformed {
+                transformed
+            } else {
+                source_kind
+            };
+            (result, transformed, fixed)
         } else {
             // If we fail to fix, lint the original source code.
             let result = lint_only(
@@ -313,8 +320,9 @@ pub(crate) fn lint_path(
                 source_type,
                 ParseSource::None,
             );
+            let transformed = source_kind;
             let fixed = FxHashMap::default();
-            (result, fixed)
+            (result, transformed, fixed)
         }
     } else {
         let result = lint_only(
@@ -326,16 +334,17 @@ pub(crate) fn lint_path(
             source_type,
             ParseSource::None,
         );
+        let transformed = source_kind;
         let fixed = FxHashMap::default();
-        (result, fixed)
+        (result, transformed, fixed)
     };
 
     let imports = imports.unwrap_or_default();
 
     if let Some((cache, relative_path, key)) = caching {
         // We don't cache parsing errors.
         if parse_error.is_none() {
-            // `FixMode::Generate` and `FixMode::Diff` rely on side-effects (writing to disk,
+            // `FixMode::Apply` and `FixMode::Diff` rely on side-effects (writing to disk,
             // and writing the diff to stdout, respectively). If a file has diagnostics, we
             // need to avoid reading from and writing to the cache in these modes.
             if match fix_mode {
@@ -350,7 +359,7 @@ pub(crate) fn lint_path(
                     LintCacheData::from_messages(
                         &messages,
                         imports.clone(),
-                        source_kind.as_ipy_notebook().map(Notebook::index).cloned(),
+                        transformed.as_ipy_notebook().map(Notebook::index).cloned(),
                     ),
                 );
             }
@@ -360,11 +369,11 @@ pub(crate) fn lint_path(
     if let Some(error) = parse_error {
         error!(
             "{}",
-            DisplayParseError::from_source_kind(error, Some(path.to_path_buf()), &source_kind,)
+            DisplayParseError::from_source_kind(error, Some(path.to_path_buf()), &transformed)
         );
     }
 
-    let notebook_indexes = if let SourceKind::IpyNotebook(notebook) = source_kind {
+    let notebook_indexes = if let SourceKind::IpyNotebook(notebook) = transformed {
         FxHashMap::from_iter([(path.to_string_lossy().to_string(), notebook.into_index())])
     } else {
         FxHashMap::default()
@@ -415,6 +424,7 @@ pub(crate) fn lint_stdin(
             data: (messages, imports),
             error: parse_error,
         },
+        transformed,
         fixed,
     ) = if matches!(fix_mode, flags::FixMode::Apply | flags::FixMode::Diff) {
         if let Ok(FixerResult {
@@ -443,8 +453,12 @@ pub(crate) fn lint_stdin(
                 }
                 flags::FixMode::Generate => {}
             }
-
-            (result, fixed)
+            let transformed = if let Cow::Owned(transformed) = transformed {
+                transformed
+            } else {
+                source_kind
+            };
+            (result, transformed, fixed)
         } else {
             // If we fail to fix, lint the original source code.
             let result = lint_only(
@@ -456,14 +470,15 @@ pub(crate) fn lint_stdin(
                 source_type,
                 ParseSource::None,
             );
-            let fixed = FxHashMap::default();
 
             // Write the contents to stdout anyway.
             if fix_mode.is_apply() {
                 source_kind.write(&mut io::stdout().lock())?;
             }
 
-            (result, fixed)
+            let transformed = source_kind;
+            let fixed = FxHashMap::default();
+            (result, transformed, fixed)
         }
     } else {
         let result = lint_only(
@@ -475,8 +490,9 @@ pub(crate) fn lint_stdin(
             source_type,
             ParseSource::None,
         );
+        let transformed = source_kind;
         let fixed = FxHashMap::default();
-        (result, fixed)
+        (result, transformed, fixed)
     };
 
     let imports = imports.unwrap_or_default();
@@ -488,7 +504,7 @@ pub(crate) fn lint_stdin(
         );
     }
 
-    let notebook_indexes = if let SourceKind::IpyNotebook(notebook) = source_kind {
+    let notebook_indexes = if let SourceKind::IpyNotebook(notebook) = transformed {
         FxHashMap::from_iter([(
             path.map_or_else(|| "-".into(), |path| path.to_string_lossy().to_string()),
             notebook.into_index(),

diff --git a/crates/ruff_cli/tests/integration_test.rs b/crates/ruff_cli/tests/integration_test.rs
@@ -284,7 +284,8 @@ fn stdin_fix_jupyter() {
    "metadata": {},
    "outputs": [],
    "source": [
-    "import os"
+    "import os\n",
+    "print(1)"
    ]
   },
   {
@@ -302,7 +303,8 @@ fn stdin_fix_jupyter() {
    "metadata": {},
    "outputs": [],
    "source": [
-    "import sys"
+    "import sys\n",
+    "print(x)"
    ]
   },
   {
@@ -354,8 +356,8 @@ fn stdin_fix_jupyter() {
  "nbformat": 4,
  "nbformat_minor": 5
 }"#), @r###"
-    success: true
-    exit_code: 0
+    success: false
+    exit_code: 1
     ----- stdout -----
     {
      "cells": [
@@ -365,7 +367,9 @@ fn stdin_fix_jupyter() {
        "id": "dccc687c-96e2-4604-b957-a8a89b5bec06",
        "metadata": {},
        "outputs": [],
-       "source": []
+       "source": [
+        "print(1)"
+       ]
       },
       {
        "cell_type": "markdown",
@@ -381,7 +385,9 @@ fn stdin_fix_jupyter() {
        "id": "cdce7b92-b0fb-4c02-86f6-e233b26fa84f",
        "metadata": {},
        "outputs": [],
-       "source": []
+       "source": [
+        "print(x)"
+       ]
       },
       {
        "cell_type": "code",
@@ -433,7 +439,8 @@ fn stdin_fix_jupyter() {
      "nbformat_minor": 5
     }
     ----- stderr -----
-    Found 2 errors (2 fixed, 0 remaining).
+    Jupyter.ipynb:cell 3:1:7: F821 Undefined name `x`
+    Found 3 errors (2 fixed, 1 remaining).
     "###);
 }
 

diff --git a/crates/ruff_linter/src/rules/pycodestyle/rules/missing_newline_at_end_of_file.rs b/crates/ruff_linter/src/rules/pycodestyle/rules/missing_newline_at_end_of_file.rs
@@ -48,9 +48,6 @@ pub(crate) fn no_newline_at_end_of_file(
     }
 
     if !source.ends_with(['\n', '\r']) {
-        // Note: if `lines.last()` is `None`, then `contents` is empty (and so we don't
-        // want to raise W292 anyway).
-        // Both locations are at the end of the file (and thus the same).
         let range = TextRange::empty(locator.contents().text_len());
 
         let mut diagnostic = Diagnostic::new(MissingNewlineAtEndOfFile, range);
@@ -60,5 +57,6 @@ pub(crate) fn no_newline_at_end_of_file(
         )));
         return Some(diagnostic);
     }
+
     None
 }