Add a SourceFile to OldDiagnostic

crates/ruff/src/commands/check.rs

@@ -131,8 +131,7 @@ pub(crate) fn check(
 
                     Diagnostics::new(
                         vec![Message::from_diagnostic(
-                            OldDiagnostic::new(IOError { message }, TextRange::default()),
-                            dummy,
+                            OldDiagnostic::new(IOError { message }, TextRange::default(), &dummy),
                             None,
                         )],
                         FxHashMap::default(),

crates/ruff/src/diagnostics.rs

@@ -69,8 +69,8 @@ impl Diagnostics {
                                     message: err.to_string(),
                                 },
                                 TextRange::default(),
+                                &source_file,
                             ),
-                            source_file,
                             None,
                         )],
                         FxHashMap::default(),
@@ -235,7 +235,7 @@ pub(crate) fn lint_path(
                     };
                     let source_file =
                         SourceFileBuilder::new(path.to_string_lossy(), contents).finish();
-                    lint_pyproject_toml(source_file, settings)
+                    lint_pyproject_toml(&source_file, settings)
                 } else {
                     vec![]
                 };
@@ -396,7 +396,7 @@ pub(crate) fn lint_stdin(
                 }
 
                 return Ok(Diagnostics {
-                    messages: lint_pyproject_toml(source_file, &settings.linter),
+                    messages: lint_pyproject_toml(&source_file, &settings.linter),
                     fixed: FixMap::from_iter([(fs::relativize_path(path), FixTable::default())]),
                     notebook_indexes: FxHashMap::default(),
                 });

crates/ruff_linter/src/checkers/ast/mod.rs

@@ -57,7 +57,7 @@ use ruff_python_semantic::{
 };
 use ruff_python_stdlib::builtins::{MAGIC_GLOBALS, python_builtins};
 use ruff_python_trivia::CommentRanges;
-use ruff_source_file::{OneIndexed, SourceRow};
+use ruff_source_file::{OneIndexed, SourceFile, SourceRow};
 use ruff_text_size::{Ranged, TextRange, TextSize};
 
 use crate::checkers::ast::annotation::AnnotationContext;
@@ -224,8 +224,6 @@ pub(crate) struct Checker<'a> {
     visit: deferred::Visit<'a>,
     /// A set of deferred nodes to be analyzed after the AST traversal (e.g., `for` loops).
     analyze: deferred::Analyze,
-    /// The cumulative set of diagnostics computed across all lint rules.
-    diagnostics: RefCell<Vec<OldDiagnostic>>,
     /// The list of names already seen by flake8-bugbear diagnostics, to avoid duplicate violations.
     flake8_bugbear_seen: RefCell<FxHashSet<TextRange>>,
     /// The end offset of the last visited statement.
@@ -239,6 +237,7 @@ pub(crate) struct Checker<'a> {
     semantic_checker: SemanticSyntaxChecker,
     /// Errors collected by the `semantic_checker`.
     semantic_errors: RefCell<Vec<SemanticSyntaxError>>,
+    collector: &'a DiagnosticsCollector<'a>,
 }
 
 impl<'a> Checker<'a> {
@@ -259,6 +258,7 @@ impl<'a> Checker<'a> {
         cell_offsets: Option<&'a CellOffsets>,
         notebook_index: Option<&'a NotebookIndex>,
         target_version: TargetVersion,
+        collector: &'a DiagnosticsCollector<'a>,
     ) -> Checker<'a> {
         let semantic = SemanticModel::new(&settings.typing_modules, path, module);
         Self {
@@ -279,7 +279,6 @@ impl<'a> Checker<'a> {
             semantic,
             visit: deferred::Visit::default(),
             analyze: deferred::Analyze::default(),
-            diagnostics: RefCell::default(),
             flake8_bugbear_seen: RefCell::default(),
             cell_offsets,
             notebook_index,
@@ -288,6 +287,7 @@ impl<'a> Checker<'a> {
             target_version,
             semantic_checker: SemanticSyntaxChecker::new(),
             semantic_errors: RefCell::default(),
+            collector,
         }
     }
 }
@@ -389,10 +389,7 @@ impl<'a> Checker<'a> {
         kind: T,
         range: TextRange,
     ) -> DiagnosticGuard<'chk, 'a> {
-        DiagnosticGuard {
-            checker: self,
-            diagnostic: Some(OldDiagnostic::new(kind, range)),
-        }
+        self.collector.report_diagnostic(kind, range)
     }
 
     /// Return a [`DiagnosticGuard`] for reporting a diagnostic if the corresponding rule is
@@ -405,15 +402,8 @@ impl<'a> Checker<'a> {
         kind: T,
         range: TextRange,
     ) -> Option<DiagnosticGuard<'chk, 'a>> {
-        let diagnostic = OldDiagnostic::new(kind, range);
-        if self.enabled(diagnostic.rule()) {
-            Some(DiagnosticGuard {
-                checker: self,
-                diagnostic: Some(diagnostic),
-            })
-        } else {
-            None
-        }
+        self.collector
+            .report_diagnostic_if_enabled(kind, range, self.settings)
     }
 
     /// Adds a [`TextRange`] to the set of ranges of variable names
@@ -2891,30 +2881,26 @@ impl<'a> Checker<'a> {
                 } else {
                     if self.semantic.global_scope().uses_star_imports() {
                         if self.enabled(Rule::UndefinedLocalWithImportStarUsage) {
-                            self.diagnostics.get_mut().push(
-                                OldDiagnostic::new(
-                                    pyflakes::rules::UndefinedLocalWithImportStarUsage {
-                                        name: name.to_string(),
-                                    },
-                                    range,
-                                )
-                                .with_parent(definition.start()),
-                            );
+                            self.report_diagnostic(
+                                pyflakes::rules::UndefinedLocalWithImportStarUsage {
+                                    name: name.to_string(),
+                                },
+                                range,
+                            )
+                            .set_parent(definition.start());
                         }
                     } else {
                         if self.enabled(Rule::UndefinedExport) {
                             if is_undefined_export_in_dunder_init_enabled(self.settings)
                                 || !self.path.ends_with("__init__.py")
                             {
-                                self.diagnostics.get_mut().push(
-                                    OldDiagnostic::new(
-                                        pyflakes::rules::UndefinedExport {
-                                            name: name.to_string(),
-                                        },
-                                        range,
-                                    )
-                                    .with_parent(definition.start()),
-                                );
+                                self.report_diagnostic(
+                                    pyflakes::rules::UndefinedExport {
+                                        name: name.to_string(),
+                                    },
+                                    range,
+                                )
+                                .set_parent(definition.start());
                             }
                         }
                     }
@@ -2975,7 +2961,8 @@ pub(crate) fn check_ast(
     cell_offsets: Option<&CellOffsets>,
     notebook_index: Option<&NotebookIndex>,
     target_version: TargetVersion,
-) -> (Vec<OldDiagnostic>, Vec<SemanticSyntaxError>) {
+    collector: &DiagnosticsCollector,
+) -> Vec<SemanticSyntaxError> {
     let module_path = package
         .map(PackageRoot::path)
         .and_then(|package| to_module_path(package, path));
@@ -3015,6 +3002,7 @@ pub(crate) fn check_ast(
         cell_offsets,
         notebook_index,
         target_version,
+        collector,
     );
     checker.bind_builtins();
 
@@ -3041,12 +3029,92 @@ pub(crate) fn check_ast(
     analyze::deferred_scopes(&checker);
 
     let Checker {
-        diagnostics,
-        semantic_errors,
-        ..
+        semantic_errors, ..
     } = checker;
 
-    (diagnostics.into_inner(), semantic_errors.into_inner())
+    semantic_errors.into_inner()
+}
+
+/// A type for collecting diagnostics in a given file.
+///
+/// [`DiagnosticsCollector::report_diagnostic`] can be used to obtain a [`DiagnosticGuard`], which
+/// will push a [`Violation`] to the contained [`OldDiagnostic`] collection on `Drop`.
+pub(crate) struct DiagnosticsCollector<'a> {
+    diagnostics: RefCell<Vec<OldDiagnostic>>,
+    source_file: &'a SourceFile,
+}
+
+impl<'a> DiagnosticsCollector<'a> {
+    /// Create a new collector with the given `source_file` and an empty collection of
+    /// `OldDiagnostic`s.
+    pub(crate) fn new(source_file: &'a SourceFile) -> Self {
+        Self {
+            diagnostics: RefCell::default(),
+            source_file,
+        }
+    }
+
+    /// Return a [`DiagnosticGuard`] for reporting a diagnostic.
+    ///
+    /// The guard derefs to an [`OldDiagnostic`], so it can be used to further modify the diagnostic
+    /// before it is added to the collection in the collector on `Drop`.
+    pub(crate) fn report_diagnostic<'chk, T: Violation>(
+        &'chk self,
+        kind: T,
+        range: TextRange,
+    ) -> DiagnosticGuard<'chk, 'a> {
+        DiagnosticGuard {
+            collector: self,
+            diagnostic: Some(OldDiagnostic::new(kind, range, self.source_file)),
+        }
+    }
+
+    /// Return a [`DiagnosticGuard`] for reporting a diagnostic if the corresponding rule is
+    /// enabled.
+    ///
+    /// Prefer [`DiagnosticsCollector::report_diagnostic`] in general because the conversion from an
+    /// `OldDiagnostic` to a `Rule` is somewhat expensive.
+    pub(crate) fn report_diagnostic_if_enabled<'chk, T: Violation>(
+        &'chk self,
+        kind: T,
+        range: TextRange,
+        settings: &LinterSettings,
+    ) -> Option<DiagnosticGuard<'chk, 'a>> {
+        let diagnostic = OldDiagnostic::new(kind, range, self.source_file);
+        if settings.rules.enabled(diagnostic.rule()) {
+            Some(DiagnosticGuard {
+                collector: self,
+                diagnostic: Some(diagnostic),
+            })
+        } else {
+            None
+        }
+    }
+
+    pub(crate) fn into_diagnostics(self) -> Vec<OldDiagnostic> {
+        self.diagnostics.into_inner()
+    }
+
+    pub(crate) fn is_empty(&self) -> bool {
+        self.diagnostics.borrow().is_empty()
+    }
+
+    pub(crate) fn retain(&self, f: impl FnMut(&OldDiagnostic) -> bool) {
+        self.diagnostics.borrow_mut().retain(f);
+    }
+
+    pub(crate) fn any(&self, f: impl FnMut(&OldDiagnostic) -> bool) -> bool {
+        self.diagnostics.borrow().iter().any(f)
+    }
+
+    pub(crate) fn swap_remove(&self, index: usize) -> OldDiagnostic {
+        self.diagnostics.borrow_mut().swap_remove(index)
+    }
+
+    /// Call `f` with an immutable borrow of the contained diagnostics.
+    pub(crate) fn with_diagnostics(&self, mut f: impl FnMut(&Vec<OldDiagnostic>)) {
+        f(&self.diagnostics.borrow());
+    }
 }
 
 /// An abstraction for mutating a diagnostic.
@@ -3058,7 +3126,7 @@ pub(crate) fn check_ast(
 /// adding fixes or parent ranges.
 pub(crate) struct DiagnosticGuard<'a, 'b> {
     /// The parent checker that will receive the diagnostic on `Drop`.
-    checker: &'a Checker<'b>,
+    collector: &'a DiagnosticsCollector<'b>,
     /// The diagnostic that we want to report.
     ///
     /// This is always `Some` until the `Drop` (or `defuse`) call.
@@ -3100,7 +3168,7 @@ impl Drop for DiagnosticGuard<'_, '_> {
         }
 
         if let Some(diagnostic) = self.diagnostic.take() {
-            self.checker.diagnostics.borrow_mut().push(diagnostic);
+            self.collector.diagnostics.borrow_mut().push(diagnostic);
         }
     }
 }

crates/ruff_linter/src/checkers/filesystem.rs

@@ -4,7 +4,7 @@ use ruff_python_ast::PythonVersion;
 use ruff_python_trivia::CommentRanges;
 
 use crate::Locator;
-use crate::OldDiagnostic;
+use crate::checkers::ast::DiagnosticsCollector;
 use crate::package::PackageRoot;
 use crate::preview::is_allow_nested_roots_enabled;
 use crate::registry::Rule;
@@ -20,40 +20,30 @@ pub(crate) fn check_file_path(
     comment_ranges: &CommentRanges,
     settings: &LinterSettings,
     target_version: PythonVersion,
-) -> Vec<OldDiagnostic> {
-    let mut diagnostics: Vec<OldDiagnostic> = vec![];
-
+    collector: &DiagnosticsCollector,
+) {
     // flake8-no-pep420
     if settings.rules.enabled(Rule::ImplicitNamespacePackage) {
         let allow_nested_roots = is_allow_nested_roots_enabled(settings);
-        if let Some(diagnostic) = implicit_namespace_package(
+        implicit_namespace_package(
             path,
             package,
             locator,
             comment_ranges,
             &settings.project_root,
             &settings.src,
             allow_nested_roots,
-        ) {
-            diagnostics.push(diagnostic);
-        }
+            collector,
+        );
     }
 
     // pep8-naming
     if settings.rules.enabled(Rule::InvalidModuleName) {
-        if let Some(diagnostic) =
-            invalid_module_name(path, package, &settings.pep8_naming.ignore_names)
-        {
-            diagnostics.push(diagnostic);
-        }
+        invalid_module_name(path, package, &settings.pep8_naming.ignore_names, collector);
     }
 
     // flake8-builtins
     if settings.rules.enabled(Rule::StdlibModuleShadowing) {
-        if let Some(diagnostic) = stdlib_module_shadowing(path, settings, target_version) {
-            diagnostics.push(diagnostic);
-        }
+        stdlib_module_shadowing(path, settings, target_version, collector);
     }
-
-    diagnostics
 }