Add TOML support to mdtest

Cargo.toml

@@ -189,6 +189,7 @@ thiserror = { version = "2.0.0" }
 thin-vec = { version = "0.2.14" }
 tikv-jemallocator = { version = "0.6.0" }
 toml = { version = "1.0.0" }
+toml_parser = { version = "1.0.0" }
 tracing = { version = "0.1.40" }
 tracing-flame = { version = "0.2.0" }
 tracing-indicatif = { version = "0.3.11" }

crates/mdtest/Cargo.toml

@@ -38,6 +38,7 @@ similar = { workspace = true }
 smallvec = { workspace = true }
 thiserror = { workspace = true }
 toml = { workspace = true }
+toml_parser = { workspace = true }
 tracing = { workspace = true }
 
 [lints]

crates/mdtest/src/assertion.rs

@@ -1,4 +1,4 @@
-//! Parse type and type-error assertions in Python comment form.
+//! Parse inline diagnostic assertions from comments.
 //!
 //! Parses comments of the form `# revealed: SomeType` and `# error: 8 [rule-code] "message text"`.
 //! In the latter case, the `8` is a column number, and `"message text"` asserts that the full
@@ -35,15 +35,21 @@
 //! ```
 
 use ruff_db::parsed::ParsedModuleRef;
-use ruff_python_ast::token::Token;
 use ruff_python_trivia::{CommentRanges, Cursor};
 use ruff_source_file::{LineIndex, OneIndexed};
 use ruff_text_size::{Ranged, TextRange, TextSize};
 use smallvec::SmallVec;
 use std::str::FromStr;
+use toml_parser::lexer::TokenKind;
 
 use crate::RunOptions;
 
+#[derive(Clone, Copy)]
+pub(crate) enum AssertionSource<'a> {
+    Python(&'a ParsedModuleRef),
+    Toml,
+}
+
 /// Diagnostic assertion comments in a single embedded file.
 #[derive(Debug)]
 pub(crate) struct InlineFileAssertions<'s> {
@@ -53,15 +59,49 @@ pub(crate) struct InlineFileAssertions<'s> {
 impl<'s> InlineFileAssertions<'s> {
     pub(crate) fn from_file(
         source: &'s str,
-        parsed: &ParsedModuleRef,
+        assertion_source: AssertionSource<'_>,
         file_index: &LineIndex,
     ) -> Self {
-        let mut by_line = Vec::new();
-        let mut file_assertions = UnparsedAssertionsIter {
-            tokens: parsed.tokens().iter(),
-            source,
+        match assertion_source {
+            AssertionSource::Python(parsed) => Self::from_comment_ranges(
+                source,
+                parsed
+                    .tokens()
+                    .iter()
+                    .filter(|token| token.kind().is_comment())
+                    .map(Ranged::range),
+                file_index,
+            ),
+            AssertionSource::Toml => Self::from_comment_ranges(
+                source,
+                toml_parser::Source::new(source)
+                    .lex()
+                    .filter(|token| token.kind() == TokenKind::Comment)
+                    .map(|token| {
+                        let span = token.span();
+                        TextRange::new(
+                            TextSize::try_from(span.start()).unwrap(),
+                            TextSize::try_from(span.end()).unwrap(),
+                        )
+                    }),
+                file_index,
+            ),
         }
-        .peekable();
+    }
+
+    fn from_comment_ranges(
+        source: &'s str,
+        comment_ranges: impl Iterator<Item = TextRange>,
+        file_index: &LineIndex,
+    ) -> Self {
+        let mut by_line = Vec::new();
+        let mut file_assertions = comment_ranges
+            .filter_map(|range| {
+                let comment_text = &source[range];
+                UnparsedAssertion::from_comment(comment_text)
+                    .map(|assertion| AssertionWithRange(assertion, range))
+            })
+            .peekable();
 
         while let Some(ranged_assertion) = file_assertions.next() {
             let mut collector = AssertionVec::new();
@@ -150,29 +190,6 @@ impl<'s> IntoIterator for InlineFileAssertions<'s> {
     }
 }
 
-struct UnparsedAssertionsIter<'a, 's> {
-    source: &'s str,
-    tokens: std::slice::Iter<'a, Token>,
-}
-
-impl<'s> Iterator for UnparsedAssertionsIter<'_, 's> {
-    type Item = AssertionWithRange<'s>;
-
-    fn next(&mut self) -> Option<Self::Item> {
-        loop {
-            let token = self.tokens.next()?;
-            if !token.kind().is_comment() {
-                continue;
-            }
-
-            let comment_text = &self.source[token.range()];
-            if let Some(assertion) = UnparsedAssertion::from_comment(comment_text) {
-                return Some(AssertionWithRange(assertion, token.range()));
-            }
-        }
-    }
-}
-
 /// An [`UnparsedAssertion`] with the [`TextRange`] of its original inline comment.
 #[derive(Debug)]
 struct AssertionWithRange<'a>(UnparsedAssertion<'a>, TextRange);
@@ -530,7 +547,18 @@ mod tests {
         db.write_file("/src/test.py", source).unwrap();
         let file = system_path_to_file(&db, "/src/test.py").unwrap();
         let parsed = parsed_module(&db, file).load(&db);
-        InlineFileAssertions::from_file(source, &parsed, &line_index(&db, file))
+        InlineFileAssertions::from_file(
+            source,
+            AssertionSource::Python(&parsed),
+            &line_index(&db, file),
+        )
+    }
+
+    fn get_toml_assertions(source: &str) -> InlineFileAssertions<'_> {
+        let mut db = TestDb::setup();
+        db.write_file("/src/ruff.toml", source).unwrap();
+        let file = system_path_to_file(&db, "/src/ruff.toml").unwrap();
+        InlineFileAssertions::from_file(source, AssertionSource::Toml, &line_index(&db, file))
     }
 
     fn into_vec(assertions: InlineFileAssertions<'_>) -> Vec<LineAssertions<'_>> {
@@ -581,6 +609,27 @@ mod tests {
         assert_eq!(format!("{assert}"), "error: ");
     }
 
+    #[test]
+    fn toml_comments() {
+        let source = dedent(
+            r##"
+            first = "# error: [not-a-comment]"
+            second = "value" # error: [rule-codes-in-selectors]
+            "##,
+        );
+        let assertions = get_toml_assertions(&source);
+
+        let [line] = &into_vec(assertions)[..] else {
+            panic!("expected one line");
+        };
+
+        assert_eq!(line.line_number, OneIndexed::from_zero_indexed(2));
+        let [assertion] = &line.assertions[..] else {
+            panic!("expected one assertion");
+        };
+        assert_eq!(format!("{assertion}"), "error: [rule-codes-in-selectors]");
+    }
+
     #[test]
     fn prior_line() {
         let source = dedent(

crates/mdtest/src/lib.rs

@@ -538,7 +538,15 @@ pub fn create_diagnostic_snapshot<'d, C>(
     writeln!(snapshot, "---").unwrap();
     writeln!(snapshot).unwrap();
 
-    writeln!(snapshot, "# Python source files").unwrap();
+    let source_heading = if test
+        .files()
+        .all(|file| matches!(file.lang, "py" | "python" | "pyi" | "ipynb"))
+    {
+        "Python source files"
+    } else {
+        "Source files"
+    };
+    writeln!(snapshot, "# {source_heading}").unwrap();
     writeln!(snapshot).unwrap();
     for file in test.files() {
         writeln!(snapshot, "## {}", file.relative_path()).unwrap();

crates/mdtest/src/matcher.rs

@@ -17,7 +17,9 @@ use ruff_source_file::{LineIndex, OneIndexed};
 use smallvec::SmallVec;
 
 use crate::RunOptions;
-use crate::assertion::{InlineFileAssertions, LineAssertions, ParsedAssertion, UnparsedAssertion};
+use crate::assertion::{
+    AssertionSource, InlineFileAssertions, LineAssertions, ParsedAssertion, UnparsedAssertion,
+};
 use crate::diagnostic::SortedDiagnostics;
 
 #[derive(Debug, Default)]
@@ -97,25 +99,34 @@ pub fn match_file(
     // Parse assertions from comments in the file, and get diagnostics from the file; both
     // ordered by line number.
     let source = source_text(db, file);
-    let parsed = parsed_module(db, file).load(db);
     let line_index = line_index(db, file);
-    let assertions = InlineFileAssertions::from_file(&source, &parsed, &line_index);
-
-    // Sort diagnostics according to the line number of the starting offset of the token in which the diagnostic appears.
-    //
-    // This can be different to the line number of the starting offset of the diagnostic range!
-    // For example, if the diagnostic is a syntax error inside a stringized annotation,
-    // the syntax error's range will likely point to a sub-range of the string literal,
-    // which will make the error unmatchable by mdtest unless we look at the token in which
-    // the diagnostic occurs (the string-literal) and use the token start as the basis for
-    // the line number.
-    let diagnostics = SortedDiagnostics::new(diagnostics, &|diagnostic_range| {
-        let token_start = parsed
-            .tokens()
-            .token_range(diagnostic_range.start())
-            .start();
-        line_index.line_index(token_start)
-    });
+    let (assertions, diagnostics) = if file.path(db).extension() == Some("toml") {
+        let assertions =
+            InlineFileAssertions::from_file(source.as_str(), AssertionSource::Toml, &line_index);
+        let diagnostics = SortedDiagnostics::new(diagnostics, &|diagnostic_range| {
+            line_index.line_index(diagnostic_range.start())
+        });
+        (assertions, diagnostics)
+    } else {
+        let parsed = parsed_module(db, file).load(db);
+        let assertions = InlineFileAssertions::from_file(
+            source.as_str(),
+            AssertionSource::Python(&parsed),
+            &line_index,
+        );
+
+        // Sort diagnostics according to the line number of the starting offset of the token in
+        // which the diagnostic appears. This can differ from the line containing the start of the
+        // diagnostic range, for example for syntax errors inside stringized annotations.
+        let diagnostics = SortedDiagnostics::new(diagnostics, &|diagnostic_range| {
+            let token_start = parsed
+                .tokens()
+                .token_range(diagnostic_range.start())
+                .start();
+            line_index.line_index(token_start)
+        });
+        (assertions, diagnostics)
+    };
 
     let mut line_diagnostics = diagnostics.iter_lines();